CPU与GPU训练模板（pytorch框架入门篇）

原始cpu训练

import torchvision
from torch.utils.data import DataLoader
from torch.nn import Sequential,Conv2d,MaxPool2d,Flatten,Linear
from torch import nn
import torch
#from mymodel import *

# 搭建网络模型====上面的数据集是十分类的网络最后的FC 一定要注意啊
from torch.utils.tensorboard import SummaryWriter

class RGD(nn.Module):
    def __init__(self):
        super(RGD, self).__init__()
        self.model = Sequential(
            Conv2d(3, 32, 5, padding=2),
            MaxPool2d(2),
            Conv2d(32, 32, 5, padding=2),
            MaxPool2d(2),
            Conv2d(32, 64, 5, padding=2),
            MaxPool2d(2),
            Flatten(),
            Linear(1024, 64),
            Linear(64, 10)
        )

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

# 数据集
train_data=torchvision.datasets.CIFAR10(root="D:\cv_box\pytorch_learning",
                                        train=True,
                                        download=True,
                                        transform=torchvision.transforms.ToTensor())

test_data=torchvision.datasets.CIFAR10(root="D:\cv_box\pytorch_learning",
                                       train=False,
                                       download=True,
                                       transform=torchvision.transforms.ToTensor())

# length长度===训练集和测试集的长度
train_data_size=len(train_data)
test_data_size=len(test_data)
print("训练数据集的长度为:{}".format(train_data_size))
print("测试集的长度为：{}".format(test_data_size))


# 数据集的加载dataloader
train_dataloader=DataLoader(train_data,batch_size=64)
test_dataloader=DataLoader(test_data,batch_size=64)


# 实例化网络的模型
rgd=RGD()

# 损失函数的目的loss=output - targets====损失函数最小值对应的自变量激素hi梯度
loss_fn=nn.CrossEntropyLoss()

# 优化器
learning_rate=0.01#learning_rate=1e-2=1x(10)^(-2)=0.01
optimizer=torch.optim.SGD(rgd.parameters(),lr=learning_rate)

# 设置网络训练的一些参数
#记录训练的次数
total_train_step=0
#记录测试的次数
total_test_step=0
#训练的轮数
epoch=10

# 使用tensorboard进行画图
writer=SummaryWriter("./accuracy_logs")

# 进入核心地带
for i in range(epoch):
    print("------------------第{}轮训练开始----------------".format(i+1))

    #训练步骤开始、
    rgd.train()#可有可无，如果有一些特殊的层（比如dropout层）必须加上，它才会调用
    for data in train_dataloader:
        imgs,targets=data
        outputs=rgd(imgs)
        loss_result=loss_fn(outputs,targets)# 损失

        # 优化器优化模型
        optimizer.zero_grad()#优化之前一定要梯度清零
        loss_result.backward()#反向传播得到当前的梯度
        optimizer.step()

        total_train_step+=1
        if total_train_step%100==0:#刻度化
            print("训练次数:{},LOSS:{}".format(total_train_step,loss_result.item()))
            writer.add_scalar("train_loss",loss_result.item(),total_train_step)

    # 测试步骤开始=====测试集就不需要调优了
    rgd.eval()#可有可无，如果有一些特殊的层（比如dropout层）必须加上，它才会调用
    total_test_loss=0
    total_accuracy=0
    with torch.no_grad():
        for data in test_dataloader:
            imgs,targets=data
            outputs=rgd(imgs)
            loss_result=loss_fn(outputs,targets)#损失，由于没有梯度，所以我们就不需要梯度优化器
            total_test_loss+=loss_result #因为loss_result是tensor的格式，需要转化为普通的数据集类型

            #图像分类里面有个独特的指标就是准确率
            accuracy= (outputs.argmax(1)==targets).sum()
            total_accuracy+=accuracy

    print("整体测试集上面的loss:{}".format(total_test_loss))#一个epoach
    print("整体测试集上的准确率acc:{}".format(total_accuracy/test_data_size))
    writer.add_scalar("test_loss",total_test_loss.item(),total_test_step)
    writer.add_scalar("test_accuracy",(total_accuracy/test_data_size).item(),total_test_step)
    total_test_step+=1#单位epoach

    #保存每一轮我们呢训练的模型结构以以及参数
    torch.save(rgd,"acc_{}.pth".format(i))
    print("模型已保存")

writer.close()
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120

GPU训练1

1、主要步骤

• 网络模型
• 损失函数
• 数据（输入、标注）
• 调用各自变量的.cuda()返回

• 对韵上面进行优化
  

GPU训练2

• 同样是修改，模型、损失函数、数据的输入和标注