Two models with same architecture, but one has an arithmetic error

Question

I am attempting to load weights from the Resnet18 backbone of an NNCLR onto a linear classifier. The issue is that I get a matrix multiplication error. In the code below, model refers to the entire NNCLR model. Much of this is based off of the Lightly documentation.

pt_backbone = model.backbone

state_dict = {'resnet18_parameters': pt_backbone.state_dict()}
torch.save(state_dict, 'test_nnclr.h5')

resnet18_new = torchvision.models.resnet18()

backbone_new = nn.Sequential(*list(resnet18_new.children())[:-1])

ckpt = torch.load('test_nnclr.h5')
backbone_new.load_state_dict(ckpt['resnet18_parameters'])
backbone_new.add_module('fc', nn.Linear(512, 2, device=cuda0))
backbone_new(torch.tensor(np.random.uniform(-10, 10, (8, 3, 128, 128)), device=cuda0).float())

I get the following runtime error:

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
 in ()
      1 backbone_new.add_module('fc', nn.Linear(512, 2, device=cuda0))
----> 2 backbone_new(torch.tensor(np.random.uniform(-10, 10, (8, 3, 128, 128)), device=cuda0).float())

4 frames
/usr/local/lib/python3.7/dist-packages/torch/nn/functional.py in linear(input, weight, bias)
   1846     if has_torch_function_variadic(input, weight, bias):
   1847         return handle_torch_function(linear, (input, weight, bias), input, weight, bias=bias)
-> 1848     return torch._C._nn.linear(input, weight, bias)
   1849 
   1850 

RuntimeError: mat1 and mat2 shapes cannot be multiplied (4096x1 and 512x2)

I understand that the 4096 comes from 512 x 8, where 8 is the batch size and 512 is the last dimension output before the linear layer. But I'm confused, because I don't see how I could account for the batch size like that in the new linear layer. I'm especially confused because of the following result:

resnet18_new(torch.tensor(np.random.uniform(-10, 10, (8, 3, 128, 128)), device=cuda0).float()).shape

Which seems to work perfectly, resulting in torch.Size([8, 2]). But the two models have the same architecture, so I don't understand how one has an error and the other doesn't. The difference between the two models is that backbone_new (which isn't actually a backbone btw) has different weights. How do I fix this error?

user_18635613 · Answer

Consider code for method forward of class ResNet. It calls attributes like self.relu or self.layer2, but right before fully-connected it calls for torch.flatten(x, 1). To understand why it's important have a look at code generated for forward of resnet18_new:

from torch.fx import symbolic_trace

symbolic_traced = symbolic_trace(resnet18_test)

print(symbolic_traced.code)

Last few lines in stdout is something like:

avgpool = self.avgpool(layer4_1_relu_1);  layer4_1_relu_1 = None
flatten = torch.flatten(avgpool, 1);  avgpool = None
fc = self.fc(flatten);  flatten = None
return fc

But similar procedure for backbone_new returns different result, after:

symbolic_traced_new = symbolic_trace(backbone_new)

print(symbolic_traced_new.code)

last lines in stdout look like (remember it's generated, so no surprise it's not pretty):

_8 = getattr(self, "8")(_7_1_relu_1);  _7_1_relu_1 = None
fc = self.fc(_8);  _8 = None
return fc

where getattr(self, "8") corresponds to AdaptiveAvgPool2d(output_size=(1, 1)). Of course it'll crash -- flatten was thrown away! And it's all because of usage torch.flatten in original ResNet._forward_impl. So to fix, when adding a new module, put there nn.Flatten also (with behavior similar to torch.flatten(x, 1)):

backbone_new = nn.Sequential(*list(resnet18_new.children())[:-1])
ckpt = torch.load('test_nnclr.h5')
backbone_new.load_state_dict(ckpt['resnet18_parameters'])
backbone_new.add_module('fc', nn.Sequential(nn.Flatten(1), nn.Linear(512, 2, device='cuda:0')))

Two models with same architecture, but one has an arithmetic error

1 Answers