Hi, I trained a model using 2 GPUs, and I want to make inference using trained model.
1. How to load this parallelised model on GPU? or multiple GPU?
2. How to load this parallelised model on CPU?
I find document mentioned the way to save the DataParallel model by add the “module”, but actually I successfully save the model in this way:
torch.save(best_model.state_dict(), bestmodel_path)
Then I use the method below to load on CPU and it works:
def load_trainedModel_cpu(model_path,device=torch.device('cpu')):
# Load trained resnet50
model = models.resnet50(pretrained=False)
num_ftrs = model.fc.in_features
model.fc = nn.Sequential(
nn.Linear(num_ftrs, 256), nn.ReLU(inplace=True), nn.Dropout(0.2),
nn.Linear(256, 3))
# Inference with CPU: use torch.load with map_location=torch.device('cpu') to map your storages to the CPU.
# original saved file with DataParallel
state_dict = torch.load(model_path,map_location=device)
# create new OrderedDict that does not contain `module.`
new_state_dict = OrderedDict()
for k, v in state_dict.items():
name = k[7:] #remove 'module'
new_state_dict[name] = v
model.load_state_dict(new_state_dict)
return model
Saving torch.nn.DataParallel Models
Save:
torch.save(model.module.state_dict(), PATH)
Load:
Load to whatever device you want
torch.nn.DataParallel is a model wrapper that enables parallel GPU utilization. To save a DataParallel model generically, save the model.module.state_dict(). This way, you have the flexibility to load the model any way you want to any device you want.
Thank you!!