pytorch-image-models

1
""" Create Conv2d Factory Method
2

3
Hacked together by / Copyright 2020 Ross Wightman
4
"""
5

6
from .mixed_conv2d import MixedConv2d
7
from .cond_conv2d import CondConv2d
8
from .conv2d_same import create_conv2d_pad
9

10

11
def create_conv2d(in_channels, out_channels, kernel_size, **kwargs):
12
    """ Select a 2d convolution implementation based on arguments
13
    Creates and returns one of torch.nn.Conv2d, Conv2dSame, MixedConv2d, or CondConv2d.
14

15
    Used extensively by EfficientNet, MobileNetv3 and related networks.
16
    """
17
    if isinstance(kernel_size, list):
18
        assert 'num_experts' not in kwargs  # MixNet + CondConv combo not supported currently
19
        if 'groups' in kwargs:
20
            groups = kwargs.pop('groups')
21
            if groups == in_channels:
22
                kwargs['depthwise'] = True
23
            else:
24
                assert groups == 1
25
        # We're going to use only lists for defining the MixedConv2d kernel groups,
26
        # ints, tuples, other iterables will continue to pass to normal conv and specify h, w.
27
        m = MixedConv2d(in_channels, out_channels, kernel_size, **kwargs)
28
    else:
29
        depthwise = kwargs.pop('depthwise', False)
30
        # for DW out_channels must be multiple of in_channels as must have out_channels % groups == 0
31
        groups = in_channels if depthwise else kwargs.pop('groups', 1)
32
        if 'num_experts' in kwargs and kwargs['num_experts'] > 0:
33
            m = CondConv2d(in_channels, out_channels, kernel_size, groups=groups, **kwargs)
34
        else:
35
            m = create_conv2d_pad(in_channels, out_channels, kernel_size, groups=groups, **kwargs)
36
    return m
37