Pytorch 1.6 多GPU训练训练Minist

最新推荐文章于 2024-03-15 08:32:15 发布

齐天大圣~~

最新推荐文章于 2024-03-15 08:32:15 发布

阅读量299

点赞数

分类专栏： Torch 文章标签：深度学习

本文链接：https://blog.csdn.net/hankerbit/article/details/108881652

版权

Torch 专栏收录该内容

2 篇文章 0 订阅

订阅专栏

这段代码演示了如何使用PyTorch进行多GPU训练，以提高处理MNIST数据集的效率。通过DataParallel将模型分布在多个GPU上，并使用Adam优化器和CrossEntropyLoss损失函数进行训练。同时，代码还展示了如何计算训练和测试阶段的准确率。

摘要由CSDN通过智能技术生成

#!/usr/bin/python3
# coding: utf-8
import torch
from torchvision import datasets, transforms
import torchvision
from tqdm import tqdm

device_ids = [1, 3, 5,6]
BATCH_SIZE = 64

# transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])])

transform = transforms.Compose([
transforms.ToTensor(),
transforms.Lambda(lambda x: x.repeat(3,1,1)),
transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
]) # RGB image

data_train = datasets.MNIST(root="./data_minist/",
transform=transform,
train=True,
download=True)
data_test = datasets.MNIST(root="./data_minist/",
transform=transform,
train=False)

data_loader_train = torch.utils.data.DataLoader(dataset=data_train,
# batch_size * num of GPU
batch_size=BATCH_SIZE * len(device_ids),
shuffle=True,
num_workers=2)

data_loader_test = torch.utils.data.DataLoader(dataset=data_test,
batch_size=BATCH_SIZE * len(device_ids),
shuffle=True,
num_workers=2)

class Model(torch.nn.Module):
def __init__(self):
super(Model, self).__init__()
self.conv1 = torch.nn.Sequential(
torch.nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1),
torch.nn.ReLU(),
torch.nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),
torch.nn.ReLU(),
torch.nn.MaxPool2d(stride=2, kernel_size=2),
)
self.dense = torch.nn.Sequential(
torch.nn.Linear(14 * 14 * 128, 1024),
torch.nn.ReLU(),
torch.nn.Dropout(p=0.5),
torch.nn.Linear(1024, 10)
)

def forward(self, x):
x = self.conv1(x)
x = x.view(-1, 14 * 14 * 128)
x = self.dense(x)
return x

model = Model()
model = torch.nn.DataParallel(model, device_ids=device_ids) # state the multi GPU
model = model.cuda(device=device_ids[0]) # put the model on the main GPU (Chose one of GPUs)

cost = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters())

n_epochs = 50
for epoch in range(n_epochs):
running_loss = 0.0
running_correct = 0
print("Epoch {}/{}".format(epoch, n_epochs))
print("-" * 10)
for data in tqdm(data_loader_train):
X_train, y_train = data
# put the data on the main GPU (choosed GPU)
X_train, y_train = X_train.cuda(device=device_ids[0]), y_train.cuda(device=device_ids[0])

outputs = model(X_train)
_, pred = torch.max(outputs.data, 1)
optimizer.zero_grad()
loss = cost(outputs, y_train)

loss.backward()
optimizer.step()
running_loss += loss.data.item()
running_correct += torch.sum(pred == y_train.data)
testing_correct = 0
for data in data_loader_test:
X_test, y_test = data
X_test, y_test = X_test.cuda(device=device_ids[0]), y_test.cuda(device=device_ids[0])
outputs = model(X_test)
_, pred = torch.max(outputs.data, 1)
testing_correct += torch.sum(pred == y_test.data)
print("Loss is:{:.4f}, Train Accuracy is:{:.4f}%, Test Accuracy is:{:.4f}".format(running_loss / len(data_train),
100 * running_correct / len(
data_train),
100 * testing_correct / len(
data_test)))
torch.save(model.state_dict(), "model_parameter.pkl")

# 使用命令行 watch -n 0.1 -d nvidia-smi 查看GPU使用率