mirror of
https://github.com/huggingface/pytorch-image-models.git
synced 2025-06-03 15:01:08 +08:00
eb7653614f
* refactor activations into basic PyTorch, jit scripted, and memory efficient custom auto * implement hard-mish, better grad for hard-swish * add initial VovNet V1/V2 impl, fix #151 * VovNet and DenseNet first models to use NormAct layers (support BatchNormAct2d, EvoNorm, InplaceIABN) * Wrap IABN for any models that use it * make more models torchscript compatible (DPN, PNasNet, Res2Net, SelecSLS) and add tests
32 lines
1.3 KiB
Python
32 lines
1.3 KiB
Python
""" Conv2d + BN + Act
|
|
|
|
Hacked together by Ross Wightman
|
|
"""
|
|
from torch import nn as nn
|
|
|
|
from .create_conv2d import create_conv2d
|
|
from .create_norm_act import convert_norm_act_type
|
|
|
|
|
|
class ConvBnAct(nn.Module):
|
|
def __init__(self, in_channels, out_channels, kernel_size=1, stride=1, padding='', dilation=1, groups=1,
|
|
norm_layer=nn.BatchNorm2d, norm_kwargs=None, act_layer=nn.ReLU, apply_act=True,
|
|
drop_block=None, aa_layer=None):
|
|
super(ConvBnAct, self).__init__()
|
|
use_aa = aa_layer is not None
|
|
self.conv = create_conv2d(
|
|
in_channels, out_channels, kernel_size, stride=1 if use_aa else stride,
|
|
padding=padding, dilation=dilation, groups=groups, bias=False)
|
|
|
|
# NOTE for backwards compatibility with models that use separate norm and act layer definitions
|
|
norm_act_layer, norm_act_args = convert_norm_act_type(norm_layer, act_layer, norm_kwargs)
|
|
self.bn = norm_act_layer(out_channels, apply_act=apply_act, drop_block=drop_block, **norm_act_args)
|
|
self.aa = aa_layer(channels=out_channels) if stride == 2 and use_aa else None
|
|
|
|
def forward(self, x):
|
|
x = self.conv(x)
|
|
x = self.bn(x)
|
|
if self.aa is not None:
|
|
x = self.aa(x)
|
|
return x
|