How to load DataParallel model which trained using multiple GPUs?

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

    return model

Saving torch.nn.DataParallel Models

Save:, PATH)


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!!

Saving the model.module.state_dict() would avoid having to manipulate the state_dict keys, but both approaches should work.

Restore the model and wrap it again into DistributedDataParallel afterwards.

Store the state_dict on the CPU or use map_location='cpu' to load the state_dict onto the CPU.