Pytorch实现手写数字识别

手写数字识别

Pytorch 实现手写数字识别，加深对网络模型理解和Pytorch代码理解

全连接网络模型构建

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
class Net_v1(nn.Module):
    # 初始化模型
    def __init__(self):
        super(Net_v1,self).__init__()
        # 全连接网络
        self.fc_layers=nn.Sequential(
            nn.Linear(1*28*28,100),
            nn.ReLU(),
            nn.Linear(100,100),
            nn.ReLU(),
            nn.Linear(100,50),
            nn.ReLU(),
            nn.Linear(50,32),
            nn.ReLU(),
            nn.Linear(32,16),
            nn.ReLU(),
            nn.Linear(16,10),
            nn.Softmax(dim=1)
        )

    def forward(self,x):
        return self.fc_layers(x)

卷积网络模型构建

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
class Net_v2(nn.Module):
    # 初始化模型
    def __init__(self):
        super(Net_v2,self).__init__()
        # 全连接网络
        self.layers=nn.Sequential(
            nn.Conv2d(1,16,(3,3)),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(16,32,(3,3)),
            nn.ReLU(),
            nn.Conv2d(32,44,(3,3)),
            nn.ReLU(),
            nn.Conv2d(44,64,(3,3)),
            nn.ReLU(),
        )

        self.outlayers = nn.Sequential(
            nn.Linear(64*7*7, 10),
            nn.Softmax(dim=1)
        )

    def forward(self,x):
        out = self.layers(x).reshape(-1,64*7*7)
        out = self.outlayers(out)
        return out

模型训练代码

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
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
class Train_v1:
    def __init__(self,weight_path):
        self.summaryWriter = SummaryWriter('./logs') # 创建文件夹存放日志

        self.net = Net_v1()  # 实例化Net_v1对象
        self.net = self.net.to(DEVICE)  # 将Net_v1对象移动到设备上
        # self.net = Net_v1.to(DEVICE) # cuda上训练

        if os.path.exists(weight_path):
            self.net.load_state_dict(torch.load(weight_path))
        
        self.opt = optim.Adam(self.net.parameters())
        self.fc_loss = nn.MSELoss()
        self.train = True
        self.test = True

    def __call__(self):
        index1,index2 = 0,0
        for epoch in range(20):
            if self.train:
                for i,(img,label) in enumerate(train_dataloader): # 自动会生成一个索引
                    #print(img.shape) # 64 1 28 28
                    #print(label.shape) # 64 
                    label = one_hot(label,10).float().to(DEVICE) # one-hot 编码，只有一个类别会显示为1，其他的为0 # to(DEVICE)放到cuda上
                    # print(label) # [0., 1., 0., 0., 0., 0., 0., 0., 0., 0.]
                    img = img.reshape(-1,1*28*28).to(DEVICE)
                    # print(img.shape) # torch.Size([64, 784])
                    # print(label.shape) # torch.Size([64, 10])
                    train_y = self.net(img) # 将图片输入网络，train_y就是预测值
                    # print(train_y.shape) # torch.Size([64, 10])
                    train_loss = self.fc_loss(train_y,label)
                
                    # 清空梯度
                    self.opt.zero_grad()

                    # 梯度计算
                    train_loss.backward()

                    # 梯度更新
                    self.opt.step()

                    if i%10==0 :
                        print(f"train_loss {i} ===>", train_loss.item())
                        self.summaryWriter.add_scalar('train_loss',train_loss,index1)
                        index1 +=1

                data_time = str(datetime.datetime.now()).replace(' ', '-').replace(':','_').replace('·','_')
                save_dir = 'param'
                if not os.path.exists(save_dir):
                    os.makedirs(save_dir)

                # 保存权重文件
                torch.save(self.net.state_dict(), f'{save_dir}/{data_time}-{epoch}.pt')
                # torch.save(self.net.state_dict(),f'param/{data_time}-{epoch}.pt')# 保存权重

            if self.test:
                # 测试
                for i,(img,label) in enumerate(test_dataloader): # 自动会生成一个索引
                 
                    label = one_hot(label,10).float().to(DEVICE) 
                    img = img.reshape(-1,1*28*28).to(DEVICE)
                
                    test_y = self.net(img) 
                    test_loss = self.fc_loss(test_y,label)
                
                    # 测试不需要梯度计算了

                    # 查看准确率
                    test_y = torch.argmax(test_y,dim=1)
                    label = torch.argmax(label,dim=1)
                    acc = torch.mean(torch.eq(test_y,label).float()) # torch.mean 均值

                    if i%10 ==0:
                        print(f"train_loss {i} ===>", test_loss.item())
                        print(f'acc {i}====>',acc.item())
                        self.summaryWriter.add_scalar('test_loss',test_loss,index2)
                        index2 +=1
                
class Train_v2:
    def __init__(self,weight_path):
        self.summaryWriter = SummaryWriter('./logs') # 创建文件夹存放日志

        self.net = Net_v2()  # 实例化Net_v1对象
        self.net = self.net.to(DEVICE)  # 将Net_v1对象移动到设备上
        # self.net = Net_v1.to(DEVICE) # cuda上训练

        if os.path.exists(weight_path):
            self.net.load_state_dict(torch.load(weight_path))
        
        self.opt = optim.Adam(self.net.parameters())
        self.fc_loss = nn.MSELoss()
        self.train = True
        self.test = True

    def __call__(self):
        index1,index2 = 0,0
        for epoch in range(20):
            if self.train:
                for i,(img,label) in enumerate(train_dataloader): # 自动会生成一个索引
                    label = one_hot(label,10).float().to(DEVICE) # one-hot 编码，只有一个类别会显示为1，其他的为0 # to(DEVICE)放到cuda上
                    # img = img.reshape(-1,1*28*28).to(DEVICE) # 图像本身就是3维的，不需要再reshape
                    # print(img.shape) # torch.Size([64, 1, 28, 28])
                    img = img.to(DEVICE)
                    train_y = self.net(img) 
                    train_loss = self.fc_loss(train_y,label)
                
                    # 清空梯度
                    self.opt.zero_grad()
                    # 梯度计算
                    train_loss.backward()
                    # 梯度更新
                    self.opt.step()

                    if i%10==0 :
                        print(f"train_loss {i} ===>", train_loss.item())
                        self.summaryWriter.add_scalar('train_loss',train_loss,index1)
                        index1 +=1

                data_time = str(datetime.datetime.now()).replace(' ', '-').replace(':','_').replace('·','_')
                save_dir = 'param'
                if not os.path.exists(save_dir):
                    os.makedirs(save_dir)

                # 保存权重文件
                torch.save(self.net.state_dict(), f'{save_dir}/{data_time}-{epoch}.pt')
                # torch.save(self.net.state_dict(),f'param/{data_time}-{epoch}.pt')# 保存权重

            if self.test:
                # 测试
                for i,(img,label) in enumerate(test_dataloader): # 自动会生成一个索引
                 
                    label = one_hot(label,10).float().to(DEVICE) 
                    # img = img.reshape(-1,1*28*28).to(DEVICE)
                    img = img.to(DEVICE)
                    test_y = self.net(img) 
                    test_loss = self.fc_loss(test_y,label)
                
                    # 测试不需要梯度计算了

                    # 查看准确率
                    test_y = torch.argmax(test_y,dim=1)
                    label = torch.argmax(label,dim=1)
                    acc = torch.mean(torch.eq(test_y,label).float()) # torch.mean 均值

                    if i%10 ==0:
                        print(f"train_loss {i} ===>", test_loss.item())
                        print(f'acc {i}====>',acc.item())
                        self.summaryWriter.add_scalar('test_loss',test_loss,index2)
                        index2 +=1

github:https://github.com/cauccliu/HandNumer

人工智能
#Python #CV
Pytorch实现手写数字识别
https://cauccliu.github.io/2024/04/09/Pytorch手写数字识别/
Author
Liuchang
Posted on
April 9, 2024
Licensed under