pytorch如何利用ResNet18进行手写数字识别

这篇文章主要介绍了pytorch如何利用ResNet18进行手写数字识别问题,具有很好的参考价值,希望对大家有所帮助。如有错误或未考虑完全的地方,望不吝赐教

利用ResNet18进行手写数字识别

先写resnet18.py

代码如下:

import torch from torch import nn from torch.nn import functional as F class ResBlk(nn.Module):     """     resnet block     """     def __init__(self, ch_in, ch_out, stride=1):         """         :param ch_in:         :param ch_out:         """         super(ResBlk, self).__init__()         self.conv1 = nn.Conv2d(ch_in, ch_out, kernel_size=3, stride=stride, padding=1)         self.bn1 = nn.BatchNorm2d(ch_out)         self.conv2 = nn.Conv2d(ch_out, ch_out, kernel_size=3, stride=1, padding=1)         self.bn2 = nn.BatchNorm2d(ch_out)         self.extra = nn.Sequential()         if ch_out != ch_in:             # [b, ch_in, h, w] => [b, ch_out, h, w]             self.extra = nn.Sequential(                 nn.Conv2d(ch_in, ch_out, kernel_size=1, stride=stride),                 nn.BatchNorm2d(ch_out)             )     def forward(self, x):         """         :param x: [b, ch, h, w]         :return:         """         out = F.relu(self.bn1(self.conv1(x)))         out = self.bn2(self.conv2(out))         # short cut         # extra module:[b, ch_in, h, w] => [b, ch_out, h, w]         # element-wise add:         out = self.extra(x) + out         out = F.relu(out)         return out class ResNet18(nn.Module):     def __init__(self):         super(ResNet18, self).__init__()         self.conv1 = nn.Sequential(             nn.Conv2d(1, 64, kernel_size=3, stride=3, padding=0),             nn.BatchNorm2d(64)         )         # followed 4 blocks         # [b, 64, h, w] => [b, 128, h, w]         self.blk1 = ResBlk(64, 128, stride=2)         # [b, 128, h, w] => [b, 256, h, w]         self.blk2 = ResBlk(128, 256, stride=2)         # [b, 256, h, w] => [b, 512, h, w]         self.blk3 = ResBlk(256, 512, stride=2)         # [b, 512, h, w] => [b, 512, h, w]         self.blk4 = ResBlk(512, 512, stride=2)         self.outlayer = nn.Linear(512 * 1 * 1, 10)     def forward(self, x):         """         :param x:         :return:         """         # [b, 1, h, w] => [b, 64, h, w]         x = F.relu(self.conv1(x))         # [b, 64, h, w] => [b, 512, h, w]         x = self.blk1(x)         x = self.blk2(x)         x = self.blk3(x)         x = self.blk4(x)         # print(x.shape) # [b, 512, 1, 1]         # 意思就是不管之前的特征图尺寸为多少,只要设置为(1,1),那么最终特征图大小都为(1,1)         # [b, 512, h, w] => [b, 512, 1, 1]         x = F.adaptive_avg_pool2d(x, [1, 1])         x = x.view(x.size(0), -1)         x = self.outlayer(x)         return x def main():     blk = ResBlk(1, 128, stride=4)     tmp = torch.randn(512, 1, 28, 28)     out = blk(tmp)     print('blk', out.shape)     model = ResNet18()     x = torch.randn(512, 1, 28, 28)     out = model(x)     print('resnet', out.shape)     print(model) if __name__ == '__main__':     main()

再写绘图utils.py

代码如下

import torch from matplotlib import pyplot as plt device = torch.device('cuda') def plot_curve(data):     fig = plt.figure()     plt.plot(range(len(data)), data, color='blue')     plt.legend(['value'], loc='upper right')     plt.xlabel('step')     plt.ylabel('value')     plt.show() def plot_image(img, label, name):     fig = plt.figure()     for i in range(6):         plt.subplot(2, 3, i + 1)         plt.tight_layout()         plt.imshow(img[i][0] * 0.3081 + 0.1307, cmap='gray', interpolation='none')         plt.title("{}: {}".format(name, label[i].item()))         plt.xticks([])         plt.yticks([])     plt.show() def one_hot(label, depth=10):     out = torch.zeros(label.size(0), depth).cuda()     idx = label.view(-1, 1)     out.scatter_(dim=1, index=idx, value=1)     return out

最后是主函数mnist_train.py

代码如下:

import torch from torch import nn from torch.nn import functional as F from torch import optim from resnet18 import ResNet18 import torchvision from matplotlib import pyplot as plt from utils import plot_image, plot_curve, one_hot batch_size = 512 # 加载数据 train_loader = torch.utils.data.DataLoader(     torchvision.datasets.MNIST('mnist_data', train=True, download=True,                                transform=torchvision.transforms.Compose([                                    torchvision.transforms.ToTensor(),                                    torchvision.transforms.Normalize(                                        (0.1307,), (0.3081,))                                ])),     batch_size=batch_size, shuffle=True) test_loader = torch.utils.data.DataLoader(     torchvision.datasets.MNIST('mnist_data/', train=False, download=True,                                transform=torchvision.transforms.Compose([                                    torchvision.transforms.ToTensor(),                                    torchvision.transforms.Normalize(                                        (0.1307,), (0.3081,))                                ])),     batch_size=batch_size, shuffle=False) # 在装载完成后,我们可以选取其中一个批次的数据进行预览 x, y = next(iter(train_loader)) # x:[512, 1, 28, 28], y:[512] print(x.shape, y.shape, x.min(), x.max()) plot_image(x, y, 'image sample') device = torch.device('cuda') net = ResNet18().to(device) optimizer = optim.SGD(net.parameters(), lr=0.01, momentum=0.9) train_loss = [] for epoch in range(5):     # 训练     net.train()     for batch_idx, (x, y) in enumerate(train_loader):         # x: [b, 1, 28, 28], y: [512]         # [b, 1, 28, 28] => [b, 10]         x, y = x.to(device), y.to(device)         out = net(x)         # [b, 10]         y_onehot = one_hot(y)         # loss = mse(out, y_onehot)         loss = F.mse_loss(out, y_onehot).to(device)         # 先给梯度清0         optimizer.zero_grad()         loss.backward()         # w' = w - lr*grad         optimizer.step()         train_loss.append(loss.item())         if batch_idx % 10 == 0:             print(epoch, batch_idx, loss.item()) plot_curve(train_loss) # we get optimal [w1, b1, w2, b2, w3, b3] # 测试 net.eval() total_correct = 0 for x, y in test_loader:     x, y = x.cuda(), y.cuda()     out = net(x)     # out: [b, 10] => pred: [b]     pred = out.argmax(dim=1)     correct = pred.eq(y).sum().float().item()     total_correct += correct total_num = len(test_loader.dataset) acc = total_correct / total_num print('test acc:', acc) x, y = next(iter(test_loader)) x, y = x.cuda(), y.cuda() out = net(x) pred = out.argmax(dim=1) x = x.cpu() pred = pred.cpu() plot_image(x, pred, 'test')

结果为:

4 90 0.009581390768289566
4 100 0.010348389856517315
4 110 0.01111914124339819
test acc: 0.9703

运行时注意把模型和参数放在GPU里,这样节省时间,此代码作为测试代码,仅供参考。

总结

以上为个人经验,希望能给大家一个参考,也希望大家多多支持0133技术站。

以上就是pytorch如何利用ResNet18进行手写数字识别的详细内容,更多请关注0133技术站其它相关文章!

赞(0) 打赏
未经允许不得转载:0133技术站首页 » python