UNet tuple stride support

src/trw/layers/blocks.py

@@ -1,5 +1,6 @@
 import collections
 import warnings
+from numbers import Number
 
 import torch
 from trw.basic_typing import TorchTensorNCX, Padding, KernelSize, Stride
@@ -227,7 +228,7 @@ def __init__(
         assert stride is not None
 
         ops = []
-        if stride != 1:
+        if (isinstance(stride, Number) and stride != 1) or (max(stride) != 1 or min(stride) != 1):
             # if stride is 1, don't upsample!
             ops.append(config.ops.upsample_fn(scale_factor=stride))
 

src/trw/layers/unet_base.py

@@ -1,4 +1,5 @@
 import copy
+from numbers import Number
 from typing import Sequence, Optional, Union, Any, List
 
 from trw.layers.convs import ModuleWithIntermediate
@@ -9,6 +10,7 @@
 from trw.utils import upsample
 from trw.layers.blocks import BlockConvNormActivation, BlockUpDeconvSkipConv, ConvBlockType
 from trw.layers.layer_config import LayerConfig, default_layer_config
+import numpy as np
 
 
 class DownType(Protocol):
@@ -248,13 +250,21 @@ def _build(self, config, init_block_fn, down_block_fn, up_block_fn, middle_block
                 out_channels = skip_channels
 
             stride = strides[len(strides) - i - 1]
+            if isinstance(stride, Number):
+                stride_minus_one = stride - 1
+            else:
+                assert len(stride) == config.ops.dim, f'expected dim={config.ops.dim}' \
+                                                      f'for `stride` but got={len(stride)}'
+                stride_minus_one = tuple(np.asarray(stride))
+                stride = tuple(stride)
+
             self.ups.append(up_block_fn(
                 config,
                 len(self.channels) - i - 1,
                 skip_channels=skip_channels,
                 input_channels=input_channels,
                 output_channels=out_channels,
-                output_padding=stride - 1,
+                output_padding=stride_minus_one,
                 stride=stride))
             input_channels = skip_channels
 

src/trw/train/outputs_trw.py

@@ -371,7 +371,7 @@ def evaluate_batch(self, batch, is_training):
         assert truth is not None, 'truth is `None` use `maybe_optional` to True'
         assert len(truth) == len(self.output), f'expected len output ({len(self.output)}) == len truth ({len(truth)})!'
         assert isinstance(truth, torch.Tensor), 'feature must be a torch.Tensor!'
-        assert truth.dtype == torch.long, 'the truth vector must be a `long` type feature'
+        assert truth.dtype == torch.long, f'the truth vector must be a `long` type feature, got={truth.dtype}'
 
         # make sure the class is not out of bound. This is a very common mistake!
         # max_index = int(torch.max(truth).cpu().numpy())