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:
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
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.