pytorch单机多卡训练的例子

用的是官网的例子，但是改成了多gpu训练。主要就是增加并行设置，同时改一下batch_size。如果你的batch_size比较大，就可以分到各个不同的gpu上。

import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor

# 指定gpu编号
device_ids = [0,1,2]

training_data = datasets.FashionMNIST(
    root="data",
    train=True,
    download=False,
    transform=ToTensor(),
)

test_data = datasets.FashionMNIST(
    root="data",
    train=False,
    download=False,
    transform=ToTensor(),
)

# 设置batch为所有的大小
batch_size = 64*len(device_ids)

# create data loaders
train_dataloader = DataLoader(training_data,batch_size=batch_size)
test_dataloader = DataLoader(test_data,batch_size=batch_size)

# device = "cuda" if torch.cuda.is_available() else "cpu"
# print(f"using {device} device")

# define a model
class NeuralNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(28*28,512),
            nn.ReLU(),
            nn.Linear(512,512),
            nn.ReLU(),
            nn.Linear(512,10)
        )

    def forward(self,x):
        x = self.flatten(x)
        logits = self.linear_relu_stack(x)
        return logits

# 设置数据并行模式
model = NeuralNetwork().to("cuda")
model = torch.nn.DataParallel(model,device_ids=device_ids)

print(model)

loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(),lr=1e-3)


def train(dataloader,model,loss_fn,optimizer):
    size = len(dataloader.dataset)
    model.train()
    for batch,(X,y) in enumerate(dataloader):
        X,y = X.to("cuda"),y.to("cuda")
        # compute prediction error
        pred = model(X)
        loss = loss_fn(pred,y)
        # backpropagation
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if batch%100==0:
            loss,current = loss.item(),batch*len(X)
            print(f"loss: {loss:.7f} [{current:>5d}/{size:>5d}]")

def test(dataloader,model,loss_fn):
    size = len(dataloader.dataset)
    num_batches = len(dataloader)
    model.eval()
    test_loss,correct = 0,0
    with torch.no_grad():
        for X,y in dataloader:
            X,y = X.to("cuda"),y.to("cuda")
            pred = model(X)
            test_loss+=loss_fn(pred,y).item()
            correct +=(pred.argmax(1)==y).type(torch.float).sum().item()

    test_loss/=num_batches
    correct/=size
    print(f"test error:\n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f}\n")


epochs = 5
for t in range(epochs):
    print(f"epoch {t+1}\n----------------------")
    train(train_dataloader,model,loss_fn,optimizer)
    test(test_dataloader,model,loss_fn)
print("Done")

torch.save(model.state_dict(),"model.pth")
print("saved pytorch model state to model.pth")