Concatenate layers with different sizes in PyTorch

Viewed 947

In Keras, it is possible to concatenate two layers of different sizes:

# Keras — this works, conceptually
layer_1 = Embedding(50, 5)(inputs)
layer_2 = Embedding(300, 20)(inputs)
concat = Concatenate()([layer_1, layer_2])
# -> `concat` now has shape `(*, 25)`, as desired

But PyTorch keeps complaining that the two layers have different sizes:

# PyTorch — this does not work
class MyModel(torch.nn.Module):
    def __init__(self):
        self.layer1 = Embedding(50, 5)
        self.layer2 = Embedding(300, 20)

    def forward(self, inputs):
        layer_1 = self.layer1(inputs)
        layer_2 = self.layer2(inputs)
        concat = torch.cat([layer_1, layer_2])

The code just above results in this error:

RuntimeError: Sizes of tensors must match except in dimension 0. Expected size 5 but got size 20 for tensor number 1 in the list.

The final concat layer I want is a layer of size 25 made of the concatenation of the two source layers.

As the two source layers are Embedding layers, I do not see as optimal that they would share the same dimension. In this example, using an embedding dimension of 5 for a vocabulary of 50 items, and an embedding dimension of size 20 for a vocabulary of 200 items.

How should this problem be solved in PyTorch?

1 Answers

Indeed torch.cat will apply the concatenation on the first axis. Since you are looking to concatenate on the second axis, you should provide the dim argument as:

>>> concat = torch.cat([layer_1, layer_2], dim=1)
Related