Repository: rwth-i6/returnn

Issue:
`TransposedConvLayer` dim tag problem
Via https://github.com/rwth-i6/pytorch-to-returnn/pull/54.

```
======================================================================
FAIL: test_layers.test_conv_transposed_2d_with_unsqueeze
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/home/runner/.local/lib/python3.8/site-packages/nose/case.py", line 198, in TestBase.runTest
    line: self.test(*self.arg)
    locals:
      self = <local> test_layers.test_conv_transposed_2d_with_unsqueeze
      self.test = <local> <function test_conv_transposed_2d_with_unsqueeze at 0x7fae50ea2790>
      self.arg = <local> ()
  File "/home/runner/work/pytorch-to-returnn/pytorch-to-returnn/tests/test_layers.py", line 239, in test_conv_transposed_2d_with_unsqueeze
    line: verify_torch_and_convert_to_returnn(model_func, inputs=x)
    locals:
      verify_torch_and_convert_to_returnn = <global> <function verify_torch_and_convert_to_returnn at 0x7fae50ea7dc0>
      model_func = <local> <function test_conv_transposed_2d_with_unsqueeze.<locals>.model_func at 0x7fae504a2f70>
      inputs = <not found>
      x = <local> array([[[ 0.49671414, -0.1382643 ,  0.64768857, ...,  0.85765964,
                           -0.15993853, -0.01901621],
                          [-1.0025294 , -0.01851314, -0.28865865, ...,  1.7547942 ,
                           -2.0819294 ,  1.6964563 ],
                          [ 0.21101747, -0.09671311, -0.5449191 , ...,  1.1487657 ,
                           -1.7397138 , -0.36..., len = 4, _[0]: {len = 512, _[0]: {len = 355}}
  File "/home/runner/work/pytorch-to-returnn/pytorch-to-returnn/pytorch_to_returnn/converter/converter.py", line 328, in verify_torch_and_convert_to_returnn
    line: converter.run()
    locals:
      converter = <local> <pytorch_to_returnn.converter.converter.Converter object at 0x7fae50426f70>
      converter.run = <local> <bound method Converter.run of <pytorch_to_returnn.converter.converter.Converter object at 0x7fae50426f70>>
  File "/home/runner/work/pytorch-to-returnn/pytorch-to-returnn/pytorch_to_returnn/converter/converter.py", line 115, in Converter.run
    line: self._run_torch_returnn_drop_in()
    locals:
      self = <local> <pytorch_to_returnn.converter.converter.Converter object at 0x7fae50426f70>
      self._run_torch_returnn_drop_in = <local> <bound method Converter._run_torch_returnn_drop_in of <pytorch_to_returnn.converter.converter.Converter object at 0x7fae50426f70>>
  File "/home/runner/work/pytorch-to-returnn/pytorch-to-returnn/pytorch_to_returnn/converter/converter.py", line 204, in Converter._run_torch_returnn_drop_in
    line: out_returnn = self._model_func(wrapped_import_torch_returnn, in_returnn)
    locals:
      out_returnn = <not found>
      self = <local> <pytorch_to_returnn.converter.converter.Converter object at 0x7fae50426f70>
      self._model_func = <local> <function test_conv_transposed_2d_with_unsqueeze.<locals>.model_func at 0x7fae504a2f70>
      wrapped_import_torch_returnn = <global> <function wrapped_import_torch_returnn at 0x7fae50ea7d30>
      in_returnn = <local> <Tensor name:? tensor:(Batch(4),512,355) returnn_data:'data' [B,F|F'feature:data'(512),T|'time:var:extern_data:data'[B]] axes id>
  File "/home/runner/work/pytorch-to-returnn/pytorch-to-returnn/tests/test_layers.py", line 235, in test_conv_transposed_2d_with_unsqueeze.<locals>.model_func
    line: return model(inputs)
    locals:
      model = <local> <ConvTranspose2d>
      inputs = <local> <Tensor name:? tensor:(Batch(4),512,355,1) returnn_data:'Unflatten_output' [B,F'feature:data'(512),T|'time:var:extern_data:data'[B],F|'Unflatten_split_dims1'(1)] axes id>
  File "/home/runner/work/pytorch-to-returnn/pytorch-to-returnn/pytorch_to_returnn/torch/nn/modules/module.py", line 414, in Module.__call__
    line: res = call_entry.apply_call()
    locals:
      res = <not found>
      call_entry = <local> <CallEntry 'ConvTranspose2d' <ModuleEntry <ConvTranspose2d>> (depth 0)>
      call_entry.apply_call = <local> <bound method CallEntry.apply_call of <CallEntry 'ConvTranspose2d' <ModuleEntry <ConvTranspose2d>> (depth 0)>>
  File "/home/runner/work/pytorch-to-returnn/pytorch-to-returnn/pytorch_to_returnn/naming/call.py", line 134, in CallEntry.apply_call
    line: layer = returnn_net.construct_layer(net_dict={layer_name: layer_dict}, name=layer_name)
    locals:
      layer = <not found>
      returnn_net = <local> <TFNetwork 'root' train=False>
      returnn_net.construct_layer = <local> <bound method TFNetwork.construct_layer of <TFNetwork 'root' train=False>>
      net_dict = <not found>
      layer_name = <local> 'ConvTranspose2d', len = 15
      layer_dict = <local> {'class': 'transposed_conv', 'from': 'Unflatten', 'activation': None, 'with_bias': True, 'n_out': 512, 'filter_size': (1, 12), 'strides': (1, 1), 'padding': 'valid', 'output_padding': (0, 0), 'remove_padding': (0, 0)}, len = 10
      name = <not found>
  File "/home/runner/.local/lib/python3.8/site-packages/returnn/tf/network.py", line 898, in TFNetwork.construct_layer
    line: return add_layer(name=name_with_prefix, layer_class=layer_class, **layer_desc)
    locals:
      add_layer = <local> <bound method TFNetwork.add_layer of <TFNetwork 'root' train=False>>
      name = <local> 'ConvTranspose2d', len = 15
      name_with_prefix = <local> 'ConvTranspose2d', len = 15
      layer_class = <local> <class 'returnn.tf.layers.basic.TransposedConvLayer'>
      layer_desc = <local> {'activation': None, 'with_bias': True, 'n_out': 512, 'filter_size': (1, 12), 'strides': (1, 1), 'padding': 'valid', 'output_padding': (0, 0), 'remove_padding': (0, 0), '_network': <TFNetwork 'root' train=False>, '_name': 'ConvTranspose2d', 'sources': [<SplitDimsLayer 'Unflatten' out_type=Data(ba..., len = 11
  File "/home/runner/.local/lib/python3.8/site-packages/returnn/tf/network.py", line 1056, in TFNetwork.add_layer
    line: layer = self._create_layer(name=name, layer_class=layer_class, **layer_desc)
    locals:
      layer = <not found>
      self = <local> <TFNetwork 'root' train=False>
      self._create_layer = <local> <bound method TFNetwork._create_layer of <TFNetwork 'root' train=False>>
      name = <local> 'ConvTranspose2d', len = 15
      layer_class = <local> <class 'returnn.tf.layers.basic.TransposedConvLayer'>
      layer_desc = <local> {'activation': None, 'with_bias': True, 'n_out': 512, 'filter_size': (1, 12), 'strides': (1, 1), 'padding': 'valid', 'output_padding': (0, 0), 'remove_padding': (0, 0), '_network': <TFNetwork 'root' train=False>, '_name': 'ConvTranspose2d', 'sources': [<SplitDimsLayer 'Unflatten' out_type=Data(ba..., len = 11
  File "/home/runner/.local/lib/python3.8/site-packages/returnn/tf/network.py", line 977, in TFNetwork._create_layer
    line: layer = layer_class(**layer_desc)
    locals:
      layer = <not found>
      layer_class = <local> <class 'returnn.tf.layers.basic.TransposedConvLayer'>
      layer_desc = <local> {'activation': None, 'with_bias': True, 'n_out': 512, 'filter_size': (1, 12), 'strides': (1, 1), 'padding': 'valid', 'output_padding': (0, 0), 'remove_padding': (0, 0), '_network': <TFNetwork 'root' train=False>, '_name': 'ConvTranspose2d', 'sources': [<SplitDimsLayer 'Unflatten' out_type=Data(ba..., len = 14
  File "/home/runner/.local/lib/python3.8/site-packages/returnn/tf/layers/basic.py", line 4373, in TransposedConvLayer.__init__
    line: output_tag.set_tag_on_size_tensor(size, batch=self.output.batch)
    locals:
      output_tag = <local> DimensionTag{'ConvTranspose2d_spatial0_transposed_conv'[?]}
      output_tag.set_tag_on_size_tensor = <local> <bound method DimensionTag.set_tag_on_size_tensor of DimensionTag{'ConvTranspose2d_spatial0_transposed_conv'[?]}>
      size = <local> <tf.Tensor 'extern_data/placeholders/data/data_dim1_size:0' shape=(?,) dtype=int32>
      batch = <not found>
      self = <local> <TransposedConvLayer 'ConvTranspose2d' out_type=Data(batch_shape_meta=[B,T|'ConvTranspose2d_spatial0_transposed_conv'[?],'ConvTranspose2d_spatial1_transposed_conv'(523),F|'Unflatten_split_dims1_replaced'(512)])>
      self.output = <local> Data(name='ConvTranspose2d_output', batch_shape_meta=[B,T|'ConvTranspose2d_spatial0_transposed_conv'[?],'ConvTranspose2d_spatial1_transposed_conv'(523),F|'Unflatten_split_dims1_replaced'(512)])
      self.output.batch = <local> BatchInfo{B}
  File "/home/runner/.local/lib/python3.8/site-packages/returnn/tf/util/data.py", line 178, in DimensionTag.set_tag_on_size_tensor
    line: assert x._is_size_of_dim_tag in (None, self)
    locals:
      x = <local> <tf.Tensor 'extern_data/placeholders/data/data_dim1_size:0' shape=(?,) dtype=int32>
      x._is_size_of_dim_tag = <local> DimensionTag{'time:var:extern_data:data'[B]}
      self = <local> DimensionTag{'ConvTranspose2d_spatial0_transposed_conv'[?]}
AssertionError: 
-------------------- >> begin captured stdout << ---------------------
>>> Running with standard reference imports...

>>> Running with wrapped imports, wrapping original PyTorch...
>>>> Module naming hierarchy:
.tmp_root: (hidden, empty)
ConvTranspose2d: <ModuleEntry ConvTranspose2d(512, 512, kernel_size=(1, 12), stride=(1, 1))> -> ...
>>>> Root module calls:
>>>> Modules with params:
>>>> Looks good!

>>> Running with wrapped Torch import, wrapping replacement for PyTorch...
RETURNN input: Data(name='data', batch_shape_meta=[B,F|F'feature:data'(512),T|'time:var:extern_data:data'[B]])
*** root/'Unflatten' layer dict: {'class': 'split_dims', 'from': 'data', 'axis': 'T', 'dims': [-1, 1]}
*** root/'Unflatten' SplitDimsLayer output: [B,F'feature:data'(512),T|'time:var:extern_data:data'[B],F|'Unflatten_split_dims1'(1)]
*** root/'ConvTranspose2d' layer dict: {'class': 'transposed_conv', 'from': 'Unflatten', 'activation': None, 'with_bias': True, 'n_out': 512, 'filter_size': (1, 12), 'strides': (1, 1), 'padding': 'valid', 'output_padding': (0, 0), 'remove_padding': (0, 0)}
Exception creating layer root/'ConvTranspose2d' of class TransposedConvLayer with opts:
{'_name': 'ConvTranspose2d',
 '_network': <TFNetwork 'root' train=False>,
 'activation': None,
 'filter_size': (1, 12),
 'n_out': 512,
 'name': 'ConvTranspose2d',
 'network': <TFNetwork 'root' train=False>,
 'output': Data(name='ConvTranspose2d_output', batch_shape_meta=[B,T|'ConvTranspose2d_spatial0_transposed_conv'[?],'ConvTranspose2d_spatial1_transposed_conv'(523),F|'Unflatten_split_dims1_replaced'(512)]),
 'output_padding': (0, 0),
 'padding': 'valid',
 'remove_padding': (0, 0),
 'sources': [<SplitDimsLayer 'Unflatten' out_type=Data(batch_shape_meta=[B,F'feature:data'(512),T|'time:var:extern_data:data'[B],F|'Unflatten_split_dims1'(1)])>],
 'strides': (1, 1),
 'with_bias': True}

--------------------- >> end captured stdout << ----------------------
```

Net dict:
```
network = {
# data: [B,F|F'feature:data'(512),T|'time:var:extern_data:data'[B]]
"Unflatten": {'class': 'split_dims', 'from': 'data', 'axis': 'T', 'dims': [-1, 1]},
# unflatten: [B,F'feature:data'(512),T|'time:var:extern_data:data'[B],F|'Unflatten_split_dims1'(1)]
"ConvTranspose2d": {
  'class': 'transposed_conv', 'from': 'Unflatten',
  'activation': None, 'with_bias': True,
  'n_out': 512, 'filter_size': (1, 12), 'strides': (1, 1),
  'padding': 'valid', 'output_padding': (0, 0), 'remove_padding': (0, 0)},
}
```

I assume this is due to the new dim tags.



Discussion:
This includes also the separate issue #596 on the split dim feature axis. But from a first glance, the `TransposedConvLayer` issue here is independent.


Resolving patch:
diff --git a/returnn/tf/layers/basic.py b/returnn/tf/layers/basic.py
--- a/returnn/tf/layers/basic.py
+++ b/returnn/tf/layers/basic.py
@@ -4640,10 +4640,12 @@ def get_out_data_from_opts(cls, name, sources, n_out,
       remove_padding = [remove_padding] * len(filter_size)
     if not isinstance(output_padding, (list, tuple)):
       output_padding = [output_padding] * len(filter_size)
-    assert len(strides) == len(out.get_spatial_axes()) == len(remove_padding) == len(output_padding), (
+    assert len(strides) == len(out.get_spatial_batch_axes()) == len(remove_padding) == len(output_padding), (
       "Expected strides for all spatial axes")
-    for idx, axis_wo_b in enumerate(out.get_spatial_axes()):
-      axis = out.get_batch_axis(axis_wo_b)
+    for idx, axis in enumerate(out.get_spatial_batch_axes()):
+      if not output_padding[idx] and remove_padding[idx] == 0 and strides[idx] == 1:
+        if filter_size[idx] == 1 or padding.lower() == "same":
+          continue
       tag = dim_tags[axis]
       dim = None
       if tag.dimension is not None:


Test patch:
diff --git a/tests/test_TFNetworkLayer.py b/tests/test_TFNetworkLayer.py
--- a/tests/test_TFNetworkLayer.py
+++ b/tests/test_TFNetworkLayer.py
@@ -7,7 +7,7 @@
 import os
 import tensorflow as tf
 import sys
-from nose.tools import assert_equal, assert_is_instance
+from nose.tools import assert_equal, assert_not_equal, assert_is_instance
 import contextlib
 import unittest
 import numpy.testing
@@ -3289,6 +3289,65 @@ def test_pool_layer_NCHW():
     print(seq_lens)
 
 
+def test_TransposedConvLayer_2d_simple():
+  # https://github.com/rwth-i6/returnn/issues/595
+  n_batch, n_time, n_in, n_out = 7, 3, 5, 13
+  config = Config({
+    "extern_data": {
+      "data": {"dim": n_in, "shape": (n_in, None)}  # [B,D,T], i.e. batch-feature-major
+    }
+  })
+  with make_scope() as session:
+    net = TFNetwork(config=config)
+    assert net.extern_data.get_default_input_data().is_batch_feature_major
+    net.construct_from_dict({
+      "Unflatten": {'class': 'split_dims', 'from': 'data', 'axis': 'T', 'dims': [-1, 1]},  # [B,D,T,1]
+      "output": {  # [B,D',T,2]
+        'class': 'transposed_conv', 'from': 'Unflatten',
+        'activation': None, 'with_bias': True,
+        'n_out': n_out, 'filter_size': (1, 2), 'strides': (1, 1),
+        'padding': 'valid', 'output_padding': (0, 0), 'remove_padding': (0, 0)},
+    })
+    out = net.get_default_output_layer().output.copy_as_batch_feature_major()
+    assert out.batch_shape == (None, 13, None, 2)
+    assert_equal(out.get_dim_tag(2), net.extern_data.get_default_input_data().get_time_dim_tag())
+    assert out.dim_tags[1].dimension == n_out and out.dim_tags[3].dimension == 2
+    in_v = numpy.arange(0, n_batch * n_time * n_in).astype("float32").reshape((n_batch, n_in, n_time))
+    session.run(tf_compat.v1.global_variables_initializer())
+    out_v = session.run(out.placeholder, feed_dict={net.extern_data.data["data"].placeholder: in_v})
+    assert isinstance(out_v, numpy.ndarray)
+    assert out_v.shape == (n_batch, n_out, n_time, 2)
+
+
+def test_TransposedConvLayer_2d_2x2():
+  n_batch, n_time, n_in, n_out = 7, 3, 5, 13
+  config = Config({
+    "extern_data": {
+      "data": {"dim": n_in, "shape": (n_in, None)}  # [B,D,T], i.e. batch-feature-major
+    }
+  })
+  with make_scope() as session:
+    net = TFNetwork(config=config)
+    assert net.extern_data.get_default_input_data().is_batch_feature_major
+    net.construct_from_dict({
+      "Unflatten": {'class': 'split_dims', 'from': 'data', 'axis': 'T', 'dims': [-1, 1]},  # [B,D,T,1]
+      "output": {  # [B,D',T,2]
+        'class': 'transposed_conv', 'from': 'Unflatten',
+        'activation': None, 'with_bias': True,
+        'n_out': n_out, 'filter_size': (2, 2), 'strides': (2, 2),
+        'padding': 'valid', 'output_padding': (0, 0), 'remove_padding': (0, 0)},
+    })
+    out = net.get_default_output_layer().output.copy_as_batch_feature_major()
+    assert out.batch_shape == (None, 13, None, 2)
+    assert_not_equal(out.get_dim_tag(2), net.extern_data.get_default_input_data().get_time_dim_tag())
+    assert out.dim_tags[1].dimension == n_out and out.dim_tags[3].dimension == 2
+    in_v = numpy.arange(0, n_batch * n_time * n_in).astype("float32").reshape((n_batch, n_in, n_time))
+    session.run(tf_compat.v1.global_variables_initializer())
+    out_v = session.run(out.placeholder, feed_dict={net.extern_data.data["data"].placeholder: in_v})
+    assert isinstance(out_v, numpy.ndarray)
+    assert out_v.shape == (n_batch, n_out, n_time * 2, 2)
+
+
 def test_ReduceLayer_NCHW():
   with make_scope() as session:
     import numpy as np