pytorch如何利用ResNet18进行手写数字识别
爱听许嵩歌 人气:0利用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里,这样节省时间,此代码作为测试代码,仅供参考。
总结
以上为个人经验,希望能给大家一个参考,也希望大家多多支持。
加载全部内容