能用兩個gpu的torch

Torch可以使用多个GPU进行训练，具体实现方法如下：

1. 使用nn.DataParallel模块：nn.DataParallel模块可以自动的在多个GPU上复制模型并进行训练，只需要在模型前加上nn.DataParallel即可。例如：

model = nn.DataParallel(model)

2. 手动划分batch到多个GPU上：手动划分batch需要使用torch.nn.parallel模块中的scatter函数，将batch数据划分到多个GPU上，然后使用parallel_apply函数在多个GPU上进行计算。例如：

import torch.nn.parallel

def train(gpu_id, model, criterion, optimizer, data_loader):
    torch.cuda.set_device(gpu_id)
    model.cuda(gpu_id)
    for batch_idx, (data, target) in enumerate(data_loader):
        data, target = data.cuda(gpu_id), target.cuda(gpu_id)
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()

model = Model()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)
train_loader = DataLoader(...)
gpu_ids = [0, 1] # 使用第0和第1块GPU
parallel = torch.nn.parallel
parallel.scatter(data, target, gpu_ids)
parallel_apply(train, gpu_ids, model, criterion, optimizer, train_loader)