added test for using unsqueeze together with transposed convolution

[LabChameleon]

I am running into trouble when using unsqueeze together with transposed convolution after the changes for the explicit data dimension tags. I dont really know what the problem is though. I added a test case which should reproduce the problem.

The error is given by:

File "/u/dierkes/pytorch-to-returnn-projects/pytorch-to-returnn-fork/pytorch-to-returnn/tests/returnn/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

[albertz]

For reference, the whole log of the failing test:

======================================================================
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 << ----------------------

[albertz]

For reference, RETURNN input:

RETURNN input: Data(name='data', batch_shape_meta=[B,F|F'feature:data'(512),T|'time:var:extern_data:data'[B]])

and RETURNN 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)},
}

[albertz]

This looks like a problem in RETURNN in TransposedConvLayer due to the new dim tag changes (https://github.com/rwth-i6/returnn/issues/595). Actually when looking at the output, it might be also another issue in RETURNN that the SplitDimsLayers marks the wrong axis as feature axis (https://github.com/rwth-i6/returnn/issues/596).

[albertz]

Note that both the RETURNN issues are fixed. So this should work now. Can you try again?

[LabChameleon]

Note that both the RETURNN issues are fixed. So this should work now. Can you try again?

Yes, it seems to work correct now.

[albertz]

Sorry, there have been some smaller changes needed to make the tests working again. Can you rebase to the latest main branch?