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
31 lines
1.3 KiB
Python
31 lines
1.3 KiB
Python
""" Create Conv2d Factory Method
|
|
|
|
Hacked together by Ross Wightman
|
|
"""
|
|
|
|
from .mixed_conv2d import MixedConv2d
|
|
from .cond_conv2d import CondConv2d
|
|
from .conv2d_same import create_conv2d_pad
|
|
|
|
|
|
def create_conv2d(in_channels, out_channels, kernel_size, **kwargs):
|
|
""" Select a 2d convolution implementation based on arguments
|
|
Creates and returns one of torch.nn.Conv2d, Conv2dSame, MixedConv2d, or CondConv2d.
|
|
|
|
Used extensively by EfficientNet, MobileNetv3 and related networks.
|
|
"""
|
|
if isinstance(kernel_size, list):
|
|
assert 'num_experts' not in kwargs # MixNet + CondConv combo not supported currently
|
|
assert 'groups' not in kwargs # MixedConv groups are defined by kernel list
|
|
# We're going to use only lists for defining the MixedConv2d kernel groups,
|
|
# ints, tuples, other iterables will continue to pass to normal conv and specify h, w.
|
|
m = MixedConv2d(in_channels, out_channels, kernel_size, **kwargs)
|
|
else:
|
|
depthwise = kwargs.pop('depthwise', False)
|
|
groups = out_channels if depthwise else kwargs.pop('groups', 1)
|
|
if 'num_experts' in kwargs and kwargs['num_experts'] > 0:
|
|
m = CondConv2d(in_channels, out_channels, kernel_size, groups=groups, **kwargs)
|
|
else:
|
|
m = create_conv2d_pad(in_channels, out_channels, kernel_size, groups=groups, **kwargs)
|
|
return m
|