pytorch基础操作(0)模型训练模板

使用CIFAR10数据集进行模型训练(模板代码)

1、准备数据集

import torch.optim
import torchvision
from torch.utils.tensorboard import SummaryWriter

from model import Net
from torch import nn
from torch.utils.data import DataLoader


# 定义训练的设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 1、准备数据集
transform = torchvision.transforms.Compose(
    [torchvision.transforms.ToTensor()]
)

train_set = torchvision.datasets.CIFAR10(root='./dataset', transform=transform, train=True, download=True)

test_set = torchvision.datasets.CIFAR10(root='./dataset', transform=transform, train=False, download=True)

train_data_size = len(train_set)
test_data_size = len(test_set)

print("训练集长度：{}".format(train_data_size))
print("测试集长度：{}".format(test_data_size))

2、加载数据集

# 2、加载数据集
train_dataloader = DataLoader(dataset=train_set,batch_size=64, shuffle=True)
test_dataloader = DataLoader(dataset=test_set,batch_size=64, shuffle=True)

3、搭建神经网络

# 3、搭建神经网络
net = Net()
# if torch.cuda.is_available():
#      net = net.cuda() # GPU训练
net = net.to(device)

import torch
from torch import nn
from torch.nn import Sequential, Conv2d, Flatten, Linear

# 搭建神经网络
class Net(nn.Module):

    def __init__(self):
        super(Net, self).__init__()
        self.model = Sequential(
            Conv2d(in_channels=3, out_channels=32, kernel_size=5, padding=2, stride=1),
            nn.MaxPool2d(kernel_size=2),
            Conv2d(in_channels=32, out_channels=32, kernel_size=5, padding=2, stride=1),
            nn.MaxPool2d(kernel_size=2),
            Conv2d(in_channels=32, out_channels=64, kernel_size=5, padding=2, stride=1),
            nn.MaxPool2d(kernel_size=2),
            Flatten(),
            Linear(in_features=1024, out_features=64),
            Linear(in_features=64, out_features=10)
        )

    def forward(self, x):
        x = self.model(x)
        return x


if __name__ == '__main__':
    net = Net()
    inputs = torch.ones(64, 3, 32, 32)
    outputs = net(inputs)
    print(outputs.shape)

4、创建loss函数

# 4、创建loss函数
loss = nn.CrossEntropyLoss()
# if torch.cuda.is_available():
#      loss = loss.cuda() # GPU训练
loss = loss.to(device)

5、创建优化器

# 5、创建优化器
learning_rate = 1e-2
optimizer = torch.optim.SGD(net.parameters(), lr=learning_rate)

6、训练网络

# 6、训练网络
# 训练的次数
train_step = 0
# 测试的次数
test_step = 0
# 训练的轮数
num_epochs = 10

# 添加tensorboard
writer = SummaryWriter('./logs')

for epoch in range(num_epochs):

    print(f' ---------- epoch {epoch + 1} ----------')
    # 训练步骤开始
    net.train()
    for imgs, targets in train_dataloader:
        # imgs = imgs.cuda()
        # targets = targets.cuda()

        train_preds = net(imgs)
        l = loss(train_preds, targets)

        # 优化器优化模型
        optimizer.zero_grad()
        l.backward()
        optimizer.step()
        train_step += 1

        if train_step % 100 == 0:
            print('训练次数: {}, loss: {}'.format(train_step, l.item()))
            writer.add_scalar('train-loss',  l.item(), train_step)

    # 测试步骤开始
    net.eval()
    sum_test_loss = 0
    total_test_acc = 0
    with torch.no_grad():
        for imgs, targets in test_dataloader:
            # imgs = imgs.cuda()
            # targets = targets.cuda()
            test_preds = net(imgs)
            l = loss(test_preds, targets)
            sum_test_loss += l
            # 正确率
            acc = ( test_preds.argmax(1) == targets).sum()
            total_test_acc += acc

    print('测试数据集上Loss: {}'.format(sum_test_loss))
    print('测试数据集上正确率acc: {}'.format(total_test_acc / test_data_size))
    writer.add_scalar('test-loss',sum_test_loss, test_step)
    writer.add_scalar('test-acc',total_test_acc / test_data_size, test_step)
    test_step += 1

    # 保存每轮模型
    torch.save(net, './model/net_{}.pth'.format(epoch))

writer.close()