<gh_stars>1-10
import torch
import torch.nn as nn

from .utils import get_same_padding
from .pooling import build_pooling_op
from lib.config import cfg


def _split_channels(num_chan, num_groups):
    split = [num_chan // num_groups for _ in range(num_groups)]
    split[0] += num_chan - sum(split)
    return split


class MixedConv2d(nn.ModuleDict):
    """ Mixed Grouped Convolution

    Based on MDConv and GroupedConv in MixNet impl:
      https://github.com/tensorflow/tpu/blob/master/models/official/mnasnet/mixnet/custom_layers.py
    """
    def __init__(self, in_channels, out_channels, kernel_size=3,
                 stride=1, dilation=1, depthwise=False, **kwargs):
        super(MixedConv2d, self).__init__()

        kernel_size = kernel_size if isinstance(kernel_size, list) else [kernel_size]
        num_groups = len(kernel_size)
        in_splits = _split_channels(in_channels, num_groups)
        out_splits = _split_channels(out_channels, num_groups)
        self.in_channels = sum(in_splits)
        self.out_channels = sum(out_splits)
        for idx, (k, in_ch, out_ch) in enumerate(zip(kernel_size, in_splits, out_splits)):
            conv_groups = out_ch if depthwise else 1
            # use add_module to keep key space clean
            padding = get_same_padding(k)
            if k <= 5:
                self.add_module(
                    str(idx),
                    nn.Conv2d(in_ch, out_ch, k, stride=stride, padding=padding, groups=conv_groups)
                )
            else:
                self.add_module(
                    str(idx),
                    build_pooling_op(cfg.model.pool, in_ch, k, padding, stride),
                )
        self.splits = in_splits

    def forward(self, x):
        x_split = torch.split(x, self.splits, 1)
        x_out = [c(x_split[i]) for i, c in enumerate(self.values())]
        x = torch.cat(x_out, 1)
        return x

import torch
import torch.nn as nn

from .utils import get_same_padding
from .pooling import build_pooling_op
from lib.config import cfg


def _split_channels(num_chan, num_groups):
    split = [num_chan // num_groups for _ in range(num_groups)]
    split[0] += num_chan - sum(split)
    return split


class MixedConv2d(nn.ModuleDict):
    """ Mixed Grouped Convolution

    Based on MDConv and GroupedConv in MixNet impl:
      https://github.com/tensorflow/tpu/blob/master/models/official/mnasnet/mixnet/custom_layers.py
    """
    def __init__(self, in_channels, out_channels, kernel_size=3,
                 stride=1, dilation=1, depthwise=False, **kwargs):
        super(MixedConv2d, self).__init__()

        kernel_size = kernel_size if isinstance(kernel_size, list) else [kernel_size]
        num_groups = len(kernel_size)
        in_splits = _split_channels(in_channels, num_groups)
        out_splits = _split_channels(out_channels, num_groups)
        self.in_channels = sum(in_splits)
        self.out_channels = sum(out_splits)
        for idx, (k, in_ch, out_ch) in enumerate(zip(kernel_size, in_splits, out_splits)):
            conv_groups = out_ch if depthwise else 1
            # use add_module to keep key space clean
            padding = get_same_padding(k)
            if k <= 5:
                self.add_module(
                    str(idx),
                    nn.Conv2d(in_ch, out_ch, k, stride=stride, padding=padding, groups=conv_groups)
                )
            else:
                self.add_module(
                    str(idx),
                    build_pooling_op(cfg.model.pool, in_ch, k, padding, stride),
                )
        self.splits = in_splits

    def forward(self, x):
        x_split = torch.split(x, self.splits, 1)
        x_out = [c(x_split[i]) for i, c in enumerate(self.values())]
        x = torch.cat(x_out, 1)
        return x
Mixed Grouped Convolution      Based on MDConv and GroupedConv in MixNet impl:       https://github.com/tensorflow/tpu/blob/master/models/official/mnasnet/mixnet/custom_layers.py # use add_module to keep key space clean