pytorch中nn.DataParallel多次使用
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
# 定义模型
class MyModel(nn.Module):
def __init__(self):
super(MyModel, self).__init__()
self.fc = nn.Linear(10, 1)
def forward(self, x):
return self.fc(x)
# 创建模型实例
model = MyModel()
# 将模型放到多张GPU上
model = nn.DataParallel(model)
model_witoout_dp = model.module