def forward(self, x: Union[Tensor,
PaddedTensor]) -> Union[Tensor, PaddedTensor]:
x, xs = (x.data, x.sizes) if isinstance(x, PaddedTensor) else (x, None)
assert x.size(1) == self.in_channels, (
f"Input image depth ({x.size(1)}) does not match the "
f"expected ({self.in_channels})")
if self.dropout and 0.0 < self.dropout < 1.0:
x = F.dropout(x, p=self.dropout, training=self.training)
x = self.conv(x)
if self.use_masks:
x = mask_image_from_size(x, batch_sizes=xs, mask_value=0)
if self.batchnorm:
x = self.batchnorm(x)
if self.activation:
x = self.activation(x)
if self.use_masks:
x = mask_image_from_size(x, batch_sizes=xs, mask_value=0)
if self.pool:
x = self.pool(x)
return (x if xs is None else PaddedTensor.build(
x, self.get_batch_output_size(xs)))
def test_padded_tensor_repr():
sizes = torch.randint(3, size=(2, 3))
t = PaddedTensor.build(torch.empty(2, 1, 5, 5), sizes)
assert (
repr(t) ==
f"PaddedTensor(data.size()=[2, 1, 5, 5], sizes={sizes.tolist()}, device=cpu)"
)
def forward(self, x):
x, xs = (x.data, x.sizes) if isinstance(x, PaddedTensor) else (x, None)
y = self._func(batch_input=x,
output_sizes=self.output_sizes,
batch_sizes=xs)
if xs is None or self._fixed_size:
return y
ys = xs.clone()
dim = int(self.output_sizes[0] is None)
ys[:, dim] = self.output_sizes[dim]
return PaddedTensor.build(y, ys)
def feed(self, x):
if isinstance(x, PaddedTensor):
x, xs = x.data, x.sizes
elif isinstance(x, tuple) and len(x) == 2:
# required because of: https://github.com/pytorch/pytorch/issues/44009
# can be deleted when we no longer support torch<=1.6.0
x, xs = x
else:
x, xs = x, None
x = self.view_as_4d(x) # N x C x H x W
if xs is not None and self._keep_padded_tensors:
if xs.size(1) == 3 and not self._keep_channels_in_size:
xs = xs[:, 1:]
return PaddedTensor.build(x, xs)
return x
def forward(self, x):
x, xs = (x.data, x.sizes) if isinstance(x, PaddedTensor) else (x, None)
assert x.size(1) == self._options.input_channels, (
f"Input image depth ({x.size(1)}) does not match "
f"the expected ({self._options.input_channels})")
x = self.conv1(x)
if self.bn1 is not None:
x = self.bn1(x)
x = self.relu(x)
x = self.maxpool(x)
x = self.layer1(x)
x = self.layer2(x)
x = self.layer3(x)
x = self.layer4(x)
if xs is None:
return x
return PaddedTensor.build(x, self.get_output_batch_size(xs))
def forward(self, x):
x, xs = (x.data, x.sizes) if isinstance(x, PaddedTensor) else (x, None)
y = self.blocks(x)
return y if xs is None else PaddedTensor.build(
y, self._compute_output_size(xs))
def test_padded_tensor_assertions(data, sizes, match):
with pytest.raises(AssertionError, match=match):
PaddedTensor.build(data, sizes)
