Pytorch DataParallel model load with map_location - pytorch

Question

Model saved with
net= Net()
model= torch.nn.DataParallel(net)
############################
# Training
############################
torch.save(model,'./model_shear_pre.pkl')
Model loading with
net = Net()
model = torch.nn.DataParallel(net, device_ids=[0,1])
model = torch.load('./model_shear_finish.pkl', map_location={'cuda:0':'cuda:0', 'cuda:1':'cuda:0', 'cuda:2':'cuda:1', 'cuda:3':'cuda:1'})
The prob is that when training I used a machine with 4 GPU, after saving the model, I would like to test it on a new machine with only 2 GPU.
After loading the saved model, I expect the model's device_ids would be [0,1], but it still be [0,1,2,3] which is the old setting. Is there anything wrong when saving or loading?

Answer 1

You should save the weights instead of the whole model.
net = Net()
model = torch.nn.DataParallel(net)
############################
# Training
############################
torch.save(model.state_dict(),'./model_shear_pre.pkl')
Then load the weight in CPU before move to all GPU
net = Net()
weights = torch.load('./model_shear_finish.pkl', map_location='cpu')
net.load_state_dict(weights)
model = torch.nn.DataParallel(net, device_ids=[0,1])
But if you have an already trained model that is saved using the whole model instead of just weights this might also work
net = torch.load('./model_shear_finish.pkl', map_location='cpu')
model = torch.nn.DataParallel(net, device_ids=[0,1])
I still recommend save only weights though. Saving and loading the whole model can really screw you up because you have to import the model the exact same way both in the save and load. And a lot of time that's a tricky thing to do. Like
train.py
from nets import Net
net = Net()
torch.save(net, './model_shear_finish.pkl')
inference.py
# this won't work
import nets
torch.load('./model_shear_finish.pkl', map_location='cpu')
# this will work
from nets import Net
torch.load('./model_shear_finish.pkl', map_location='cpu')