训练:
'''
定义device,其中需要注意的是“cuda:0”代表起始的device_id为0,如果直接是“cuda”,
同样默认是从0开始。可以根据实际需要修改起始位置,如“cuda:1”。
'''
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = Model(input_size, output_size) # 实例化模型对象
# 多张GPU同时使用
if torch.cuda.device_count() > 1: # 检查电脑是否有多块GPU
print(f"Let's use {torch.cuda.device_count()} GPUs!")
model = nn.DataParallel(model) # 将模型对象转变为多GPU并行运算的模型
# model = nn.DataParallel(model,device_ids=[0,1,2])
model.to(device) # 把并行的模型移动到GPU上
测试:先并行,再加载模型
net=resnet50(pretrained=True)
in_channel=net.fc.in_features
net.fc = nn.Linear(in_channel, 10)
if torch.cuda.device_count() > 1:
print(f"Let's use {torch.cuda.device_count()} GPUs!")
net = nn.DataParallel(net)
net.load_state_dict(torch.load("resnet50state.pth"), strict=False)
引用: