How to switch model trained on 2 GPUs to 1 GPU?

4

The problem here is, that you have saved your model as torch.DataParallel just wrap it in the DataParallel(model) before you start training and and specify the max number of GPUs to use as a workaround. The more elegant mehtod would be to change the saved state_dict.
The problem is "module.main.0.0.weight" the "module." of each key in your saved state dict, just remove it and save it again (be sure to specifiy the map_location argument in torch.load() to only one GPU: Here is the code to remove the "module." string before all the keys:

state_dict = torch.load('path/to/weights')
keys = state_dict.keys()
values = state_dict.values()

new_keys = []
for key in keys:
    new_key = key[8:]    # remove the 'module.'
    new_keys.append(new_key)

new_dict = OrderedDict(list(zip(new_keys, values))) # create a new OrderedDict with (key, value) pairs
torch.save(new_dict, 'path/to/save')
2 Likes