A wrapper around :class:`torch.nn.Conv2d` to support empty inputs and more features.
| 5 | from collections import namedtuple |
| 6 | |
| 7 | class Conv2d(torch.nn.Conv2d): |
| 8 | """ |
| 9 | A wrapper around :class:`torch.nn.Conv2d` to support empty inputs and more features. |
| 10 | """ |
| 11 | |
| 12 | def __init__(self, *args, **kwargs): |
| 13 | """ |
| 14 | Extra keyword arguments supported in addition to those in `torch.nn.Conv2d`: |
| 15 | |
| 16 | Args: |
| 17 | norm (nn.Module, optional): a normalization layer |
| 18 | activation (callable(Tensor) -> Tensor): a callable activation function |
| 19 | |
| 20 | It assumes that norm layer is used before activation. |
| 21 | """ |
| 22 | norm = kwargs.pop("norm", None) |
| 23 | activation = kwargs.pop("activation", None) |
| 24 | super().__init__(*args, **kwargs) |
| 25 | |
| 26 | self.norm = norm |
| 27 | self.activation = activation |
| 28 | |
| 29 | def forward(self, x): |
| 30 | # torchscript does not support SyncBatchNorm yet |
| 31 | # https://github.com/pytorch/pytorch/issues/40507 |
| 32 | # and we skip these codes in torchscript since: |
| 33 | # 1. currently we only support torchscript in evaluation mode |
| 34 | # 2. features needed by exporting module to torchscript are added in PyTorch 1.6 or |
| 35 | # later version, `Conv2d` in these PyTorch versions has already supported empty inputs. |
| 36 | if not torch.jit.is_scripting(): |
| 37 | if x.numel() == 0 and self.training: |
| 38 | # https://github.com/pytorch/pytorch/issues/12013 |
| 39 | assert not isinstance( |
| 40 | self.norm, torch.nn.SyncBatchNorm |
| 41 | ), "SyncBatchNorm does not support empty inputs!" |
| 42 | |
| 43 | x = F.conv2d( |
| 44 | x, self.weight, self.bias, self.stride, self.padding, self.dilation, self.groups |
| 45 | ) |
| 46 | if self.norm is not None: |
| 47 | x = self.norm(x) |
| 48 | if self.activation is not None: |
| 49 | x = self.activation(x) |
| 50 | return x |
| 51 | |
| 52 | |
| 53 | class ShapeSpec(namedtuple("_ShapeSpec", ["channels", "height", "width", "stride"])): |