import torch
import torch.nn as nn
import torch.optim as optim
model = Model()
device_ids = [0, 1]
model = nn.DataParallel(model, device_ids=device_ids)
model.to(device_ids[0]) # 选择第一个GPU作为主GPU
optimizer = optim.SGD(model.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss()
# 准备输入数据和标签(需要在GPU上)
inputs = inputs.to(device_ids[0])
labels = labels.to(device_ids[0])
单机多卡训练
最新推荐文章于 2024-09-14 09:59:20 发布