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
| 78 | |
| 79 | |
| 80 | class MixedConv2d(nn.Module): |
| 81 | """ Mixed Grouped Convolution |
| 82 | Based on MDConv and GroupedConv in MixNet impl: |
| 83 | https://github.com/tensorflow/tpu/blob/master/models/official/mnasnet/mixnet/custom_layers.py |
| 84 | """ |
| 85 | |
| 86 | def __init__(self, in_channels, out_channels, kernel_size=3, |
| 87 | stride=1, padding='', split_op='equal', dilated=False, depthwise=False, **kwargs): |
| 88 | super(MixedConv2d, self).__init__() |
| 89 | |
| 90 | kernel_size = kernel_size if isinstance(kernel_size, list) else [kernel_size] |
| 91 | num_groups = len(kernel_size) |
| 92 | in_splits = _split_channels(in_channels, num_groups, split_op=split_op) |
| 93 | out_splits = _split_channels(out_channels, num_groups, split_op=split_op) |
| 94 | for idx, (k, in_ch, out_ch) in enumerate(zip(kernel_size, in_splits, out_splits)): |
| 95 | d = 1 |
| 96 | # FIXME make compat with non-square kernel/dilations/strides |
| 97 | if stride == 1 and dilated: |
| 98 | d, k = (k - 1) // 2, 3 |
| 99 | conv_groups = out_ch if depthwise else 1 |
| 100 | # use add_module to keep key space clean |
| 101 | self.add_module( |
| 102 | str(idx), |
| 103 | conv2d_pad( |
| 104 | in_ch, out_ch, k, stride=stride, |
| 105 | padding=padding, dilation=d, groups=conv_groups, **kwargs) |
| 106 | ) |
| 107 | self.splits = in_splits |
| 108 | |
| 109 | def forward(self, x): |
| 110 | if len(self.splits) > 1: |
| 111 | x_split = torch.split(x, self.splits, 1) |
| 112 | x_out = [c(x) for x, c in zip(x_split, self._modules.values())] |
| 113 | x = torch.cat(x_out, 1) |
| 114 | else: |
| 115 | x_out = [c(x) for c in self._modules.values()] |
| 116 | x = x_out[0] |
| 117 | return x |
| 118 | |
| 119 | |
| 120 | # helper method |