[Python] Why dimension errors are reported when the dimensions are actually correct

I split the model and gave them to different clients for training, and used mpi4py to pass parameters. The following is the error when I ran 3 processes. Note that the output shape in the printout refers to the output dimension of the previous layer, that is, the input dimension of this layer.

Epoch: 0
inputs shape: torch.Size([102, 3, 32, 32]) targets shape: torch.Size([102])
Client 0 processing layer conv1
output shape: torch.Size([102, 3, 32, 32])
Client 0 processing layer bn1
output shape: torch.Size([102, 64, 32, 32])
Client 0 processing layer layer1.0.bn1
output shape: torch.Size([102, 64, 32, 32])
Client 0 processing layer layer1.0.conv1
output shape: torch.Size([102, 64, 32, 32])
Client 0 processing layer layer1.0.bn2
output shape: torch.Size([102, 64, 32, 32])
Client 0 processing layer layer1.0.conv2
output shape: torch.Size([102, 64, 32, 32])
Client 0 processing layer layer1.1.bn1
output shape: torch.Size([102, 64, 32, 32])
Client 0 processing layer layer1.1.conv1
output shape: torch.Size([102, 64, 32, 32])
Client 0 processing layer layer1.1.bn2
output shape: torch.Size([102, 64, 32, 32])
Client 0 processing layer layer1.1.conv2
output shape: torch.Size([102, 64, 32, 32])
Client 0 processing layer layer2.0.bn1
output shape: torch.Size([102, 64, 32, 32])
Client 0 processing layer layer2.0.conv1
output shape: torch.Size([102, 64, 32, 32])
inputs shape: torch.Size([102, 3, 32, 32]) targets shape: torch.Size([102])
Client 0 processing layer conv1
output shape: torch.Size([102, 3, 32, 32])
train_split: Client 1 received prev_output
Client 0 processing layer bn1
output shape: torch.Size([102, 64, 32, 32])
Client 0 processing layer layer1.0.bn1
output shape: torch.Size([102, 64, 32, 32])
Client 0 processing layer layer1.0.conv1
output shape: torch.Size([102, 64, 32, 32])
Traceback (most recent call last):
File "/home/AI_blockchain/watermark-split/train.py", line 125, in <module>
train_split(epoch, net, criterion, optimizer, logfile,
File "/home/AI_blockchain/watermark-split/trainer.py", line 132, in train_split
output = layer(output)
^^^^^^^^^^^^^
File "/home/anaconda3/envs/wmsplit/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/anaconda3/envs/wmsplit/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/anaconda3/envs/wmsplit/lib/python3.11/site-packages/torch/nn/modules/conv.py", line 460, in forward
return self._conv_forward(input, self.weight, self.bias)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/anaconda3/envs/wmsplit/lib/python3.11/site-packages/torch/nn/modules/conv.py", line 456, in _conv_forward
return F.conv2d(input, weight, bias, self.stride,
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: Given groups=1, weight of size [128, 64, 1, 1], expected input[102, 128, 16, 16] to have 64 channels, but got 128 channels instead
Client 0 processing layer layer1.0.bn2
output shape: torch.Size([102, 64, 32, 32])
Client 0 processing layer layer1.0.conv2
output shape: torch.Size([102, 64, 32, 32])
Client 0 processing layer layer1.1.bn1
output shape: torch.Size([102, 64, 32, 32])
Client 0 processing layer layer1.1.conv1
output shape: torch.Size([102, 64, 32, 32])
Client 0 processing layer layer1.1.bn2
output shape: torch.Size([102, 64, 32, 32])
Client 0 processing layer layer1.1.conv2
output shape: torch.Size([102, 64, 32, 32])
Client 0 processing layer layer2.0.bn1
output shape: torch.Size([102, 64, 32, 32])
Client 0 processing layer layer2.0.conv1
output shape: torch.Size([102, 64, 32, 32])
train_split: Client 2


As you can see, when the client 0 processes layer1.0.conv1 for the first time, there is no error, and when the client 0 processes layer1.0.conv1 for the second time, there is a RuntimeError: Given groups=1, weight of size [128, 64, 1, 1], expected input[102, 128, 16, 16] to have 64 channels, but got 128 channels instead. But as we can see from the dimensions printed above, layer1.0.conv1's input dimensions are [102, 64, 32, 32], the number of channels is 64 instead of 128 as indicated in the error message. Run to print train_split: Client 2 and it's stuck. Thank you in advance for your answer!

Continue reading...