pytorch实现lenet5

import gzip import struct import numpy as np import torch import torch.nn as nn import torch.nn.functional as F import torch.optim as optim # 读取数据的函数,先读取标签，再读取图片 def _read(image, label): with gzip.open(label) as flbl: magic, num = struct.unpack(">II", flbl.read(8)) label = np.frombuffer(flbl.read(), dtype=np.int8) with gzip.open(image, 'rb') as fimg: magic, num, rows, cols = struct.unpack(">IIII", fimg.read(16)) image = np.frombuffer(fimg.read(), dtype=np.uint8) image = image.reshape(len(label), rows, cols) return image, label # 读取数据 def get_data(): train_img, train_label = _read( 'train-images-idx3-ubyte.gz', 'train-labels-idx1-ubyte.gz') test_img, test_label = _read( 't10k-images-idx3-ubyte.gz', 't10k-labels-idx1-ubyte.gz') return [train_img, train_label, test_img, test_label] class Network(nn.Module): def __init__(self): super().__init__() self.maxpool = nn.MaxPool2d(kernel_size=2, stride=2) self.layer1conv = nn.Conv2d(in_channels=1, out_channels=6, kernel_size=5, padding=2) self.layer2conv = nn.Conv2d(in_channels=6, out_channels=16, kernel_size=5) self.layer3conv = nn.Conv2d(in_channels=16, out_channels=120, kernel_size=5) self.layer4line = nn.Linear(in_features=120, out_features=84) self.layer5line = nn.Linear(in_features=84, out_features=10) def forward(self, x): x = F.relu(self.layer1conv(x)) # 6*28*28 x = self.maxpool(x) # 6*14*14 x = F.relu(self.layer2conv(x)) # 16*10*10 x = self.maxpool(x) # 16*5*5 x = F.relu(self.layer3conv(x)) # 120*1*1 x = x.view(x.size(0), -1) # 120 x = F.relu(self.layer4line(x)) # 84 x = self.layer5line(x) # 10 return x network = Network() network.load_state_dict(torch.load('network_data.pth')) train_set_x_orig, train_set_y_orig, test_set_x_orig, test_set_y_orig = get_data() train_set_x_orig = torch.tensor(train_set_x_orig, dtype=torch.float32)/128-1 test_set_x_orig = torch.tensor(test_set_x_orig, dtype=torch.float32)/128-1 train_set_y_orig = torch.tensor(train_set_y_orig, dtype=torch.int64) test_set_y_orig = torch.tensor(test_set_y_orig, dtype=torch.int64) test_x = test_set_x_orig.unsqueeze_(1) pred_test = network(test_x) maxcorrect = pred_test.argmax(dim=1).eq(test_set_y_orig).sum().item() optimizer = optim.Adam(network.parameters(), lr=0.001) for t in range(600): train_x = train_set_x_orig[100*t:100*t+100].unsqueeze_(1) train_y = train_set_y_orig[100*t:100*t+100] pred_train = network(train_x) train_loss = F.cross_entropy(pred_train, train_y) optimizer.zero_grad() train_loss.backward() optimizer.step() print(train_loss.item()) pred_test = network(test_x) correct = pred_test.argmax(dim=1).eq(test_set_y_orig).sum().item() if correct > maxcorrect: torch.save(network.state_dict(), 'network_data.pth') maxcorrect = correct