Try native support for empty inputs
Summary:
Transition to native pytorch support that's recently added:
adds support for conv/convtranspose.
adds support for batchnorm.
adds support for groupnorm.
@@ -12,6 +12,8 @@ import math
import torch
from torch.nn.modules.utils import _ntuple
+TORCH_VERSION = tuple(int(x) for x in torch.__version__.split(".")[:2])
+
def cat(tensors, dim=0):
"""
@@ -58,7 +60,16 @@ class Conv2d(torch.nn.Conv2d):
self.activation = activation
def forward(self, x):
- if x.numel() == 0:
+ if x.numel() == 0 and self.training:
+ # https://github.com/pytorch/pytorch/issues/12013
+ assert not isinstance(
+ self.norm, torch.nn.SyncBatchNorm
+ ), "SyncBatchNorm does not support empty inputs!"
+
+ if x.numel() == 0 and TORCH_VERSION <= (1, 4):
+ assert not isinstance(
+ self.norm, torch.nn.GroupNorm
+ ), "GroupNorm does not support empty inputs in PyTorch <=1.4!"
# When input is empty, we want to return a empty tensor with "correct" shape,
# So that the following operations will not panic
# if they check for the shape of the tensor.
@@ -72,10 +83,6 @@ class Conv2d(torch.nn.Conv2d):
output_shape = [x.shape[0], self.weight.shape[0]] + output_shape
empty = _NewEmptyTensorOp.apply(x, output_shape)
if self.training:
- # https://github.com/pytorch/pytorch/issues/12013
- assert not isinstance(
- self.norm, torch.nn.SyncBatchNorm
- ), "SyncBatchNorm does not support empty inputs!"
# This is to make DDP happy.
# DDP expects all workers to have gradient w.r.t the same set of parameters.
_dummy = sum(x.view(-1)[0] for x in self.parameters()) * 0.0
@@ -91,6 +98,10 @@ class Conv2d(torch.nn.Conv2d):
return x
+if TORCH_VERSION > (1, 4):
+ ConvTranspose2d = torch.nn.ConvTranspose2d
+else:
+
class ConvTranspose2d(torch.nn.ConvTranspose2d):
"""
A wrapper around :class:`torch.nn.ConvTranspose2d` to support zero-size tensor.
@@ -123,6 +134,10 @@ class ConvTranspose2d(torch.nn.ConvTranspose2d):
return _NewEmptyTensorOp.apply(x, output_shape) + _dummy
+if TORCH_VERSION > (1, 4):
+ BatchNorm2d = torch.nn.BatchNorm2d
+else:
+
class BatchNorm2d(torch.nn.BatchNorm2d):
"""
A wrapper around :class:`torch.nn.BatchNorm2d` to support zero-size tensor.
