Loading a checkpoint (stored as a “pth.tar”) into DenseNet to then later convert it to Keras. The code below works but gives issues with formatting during conversion later. First, I simply loaded the state dict from the “pth.tar” without changing classifier weight and bias tensor shapes but was getting torch.size tensor mismatch. I’ve initialized those required tensor shapes using the data attribute. Am I not supposed to be messing with the “.data” attribute? Basically, what do I need to change in my code so I can load such “pth.tar” checkpoints and still have the model work for downstream tasks?
new_state_dict = OrderedDict()
model = models.densenet121(pretrained=True)
# to load
model.classifier.weight.data = torch.randn(42, 1024)
model.classifier.bias.data = torch.randn(42)
checkpoint = torch.load('../iter_196800.pth.tar')
for k, v in checkpoint['model_state'].items():
if k[:13] == 'module.model.':
name = k[13:] # remove `module.model`
else:
name = k
new_state_dict[name] = v
checkpoint['model_state'] = new_state_dict
model.load_state_dict(checkpoint['model_state'])
#optimizer.load_state_dict(checkpoint['optimizer'])
model.eval()