这块之前已经操练过了，把网络结构与训练函数放在这里。

网络结构

  1. import time
  2. import torch
  3. from torch import nn, optim
  4. import torchvision
  5. from utils import *
  8. class LeNet(nn.Module):
  9.     def __init__(self):
  10.         super(LeNet, self).__init__()
  11.         self.conv = nn.Sequential(
  12.             nn.Conv2d(1, 6, 5), # in_channels, out_channels, kernel_size
  13.             nn.ReLU(),
  14.             nn.MaxPool2d(2, 2), # kernel_size, stride
  15.             nn.Conv2d(6, 16, 5),
  16.             nn.ReLU(),
  17.             nn.MaxPool2d(2, 2)
  18.         )
  19.         self.fc = nn.Sequential(
  20.             nn.Linear(16*4*4, 512),
  21.             nn.LeakyReLU(negative_slope=0.01),
  22.             nn.Dropout(0.5),
  23.             # nn.ReLU(),
  24.             nn.Linear(512, 256),
  25.             # nn.LeakyReLU(negative_slope=0.01),
  26.             nn.Sigmoid(),
  27.             nn.Dropout(0.5),
  28.             nn.Linear(256, 200)
  29.         )
  31.     def forward(self, img):
  32.         feature = self.conv(img)
  33.         output = self.fc(feature.view(img.shape[0], -1))
  34.         return output
  37. if __name__ == "__main__":
  38.     net = LeNet()
  39.     print(net)
  40.     batch_size = 64
  41.     resize = [28, 28]
  42.     train_iter, test_iter = load_data_cnn(batch_size, resize=resize)
  43.     lr, num_epochs = 0.001, 50
  44.     loss = torch.nn.CrossEntropyLoss()
  45.     optimizer = torch.optim.Adam(net.parameters(), lr=lr, weight_decay=0.001)
  46.     train(net, train_iter, test_iter, loss=loss, optimizer=optimizer, device='cuda', num_epochs=num_epochs, file='cnn.pt')

数据集与训练函数

  1. import torch
  2. from torch.utils.data import *
  3. import struct
  4. import os
  5. import time
  6. import torchvision
  7. from PIL import Image
  8. from torchvision.datasets.vision import VisionDataset
  9. from torchvision.transforms import ToPILImage
  10. # from IPython.display import Image
  11. import matplotlib.pyplot as plt
  12. from IPython import display
  13. import hiddenlayer as hl
  16. class ImageSet(Dataset):
  17.     def __init__(self, path, dimensions, classes, size):
  18.         super(ImageSet, self).__init__()
  19.         data_raw = ()
  20.         label_raw = []
  21.         for i in range(classes):
  22.             file_name = os.path.join(path, 'f' + str(i) + '.dat')
  23.             with open(file_name, 'rb') as f:
  24.                 data_raw += struct.unpack('f' * size * dimensions, f.read(4 * size * dimensions))
  25.                 for j in range(size):
  26.                     label_raw.append(float(i))
  27.         self.train_data = torch.tensor(data_raw).clone().view(-1, dimensions)
  28.         self.train_label = torch.tensor(label_raw)
  30.     def __getitem__(self, index):
  31.         return self.train_data[index], self.train_label[index].long()
  33.     def __len__(self):
  34.         return self.train_label.size()[0]
  37. class ImageSetCNN(VisionDataset):
  38.     def __init__(self, root, dim_x, dim_y, classes, size, transform=None, target_transform=None):
  39.         super(ImageSetCNN, self).__init__(root, transform=transform, target_transform=target_transform)
  40.         data_raw = ()
  41.         label_raw = []
  42.         for i in range(classes):
  43.             file_name = os.path.join(root, 'f' + str(i) + '.dat')
  44.             with open(file_name, 'rb') as f:
  45.                 data_raw += struct.unpack('f' * size * dim_x * dim_y, f.read(4 * size * dim_x * dim_y))
  46.                 for j in range(size):
  47.                     # one_hot = torch.zeros(classes)
  48.                     # one_hot[i] = 1
  49.                     label_raw.append(i)
  50.         self.train_data = torch.tensor(data_raw).clone().view(-1, dim_x, dim_y)
  51.         self.train_label = torch.tensor(label_raw)
  53.     def __getitem__(self, index):
  54.         img, target = self.train_data[index], int(self.train_label[index])
  56.         img = Image.fromarray(img.numpy(), mode='L')
  57.         # img.show()
  59.         if self.transform is not None:
  60.             img = self.transform(img)
  62.         if self.target_transform is not None:
  63.             target = self.target_transform(target)
  65.         return img, target
  67.     def __len__(self):
  68.         return self.train_label.size()[0]
  71. def load_data(batch_size):
  72.     training_set = ImageSet('train', 440, 200, 144)
  73.     test_set = ImageSet('test', 440, 200, 18)
  74.     num_workers = 4
  75.     train_iter = DataLoader(training_set, batch_size=batch_size, shuffle=True, num_workers=num_workers)
  76.     test_iter = DataLoader(test_set, batch_size=batch_size, shuffle=True, num_workers=num_workers)
  77.     return train_iter, test_iter
  80. def load_data_cnn(batch_size, resize=None):
  81.     trans = []
  82.     if resize:
  83.         trans.append(torchvision.transforms.Resize(size=resize))
  84.     trans.append(torchvision.transforms.ToTensor())
  86.     transform = torchvision.transforms.Compose(trans)
  87.     training_set = ImageSetCNN('train', 20, 22, 200, 144, transform=transform)
  88.     test_set = ImageSetCNN('test', 20, 22, 200, 18, transform=transform)
  90.     train_iter = DataLoader(training_set, batch_size=batch_size, shuffle=True, num_workers=4)
  91.     test_iter = DataLoader(test_set, batch_size=batch_size, shuffle=False, num_workers=4)
  93.     return train_iter, test_iter
  96. def train(net, train_iter, test_iter, loss, optimizer, device, num_epochs, file):
  97.     # 记录训练过程的指标
  98.     history = hl.History()
  99.     # 使用canvas进行可视化
  100.     canvas = hl.Canvas()
  101.     net = net.to(device)
  102.     print("training on ", device)
  103.     best_acc = 0
  104.     no_progress = 0
  105.     # loss = torch.nn.CrossEntropyLoss()
  106.     for epoch in range(num_epochs):
  107.         train_l_sum, train_acc_sum, n, batch_count, start = 0.0, 0.0, 0, 0, time.time()
  108.         for X, y in train_iter:
  109.             X = X.to(device)
  110.             y = y.to(device)
  111.             y_hat = net(X)
  112.             l = loss(y_hat, y)
  113.             optimizer.zero_grad()
  114.             l.backward()
  115.             optimizer.step()
  116.             train_l_sum += l.cpu().item()
  117.             train_acc_sum += (y_hat.argmax(dim=1) == y).sum().cpu().item()
  118.             n += y.shape[0]
  119.             batch_count += 1
  120.         test_acc = evaluate(test_iter, net)
  121.         print('epoch %d, loss %.4f, train acc %.3f, test acc %.3f, time %.1f sec'
  122.               % (epoch + 1, train_l_sum / batch_count, train_acc_sum / n, test_acc, time.time() - start))
  123.         history.log(epoch,
  124.                     train_loss=train_l_sum / batch_count,
  125.                     train_acc=train_acc_sum / n,
  126.                     test_acc=test_acc)
  127.         # 可视化
  128.         with canvas:
  129.             canvas.draw_plot(history["train_loss"])
  130.             canvas.draw_plot(history["train_acc"])
  131.             canvas.draw_plot(history["test_acc"])
  132.         if test_acc > best_acc:
  133.             no_progress = 0
  134.             torch.save(net.state_dict(), file)
  135.             best_acc = test_acc
  136.             print('best model saved to ' + file)
  137.         else:
  138.             no_progress += 1
  139.             if no_progress > 20:
  140.                 break
  143. def evaluate(data_iter, net, device=None):
  144.     if device is None and isinstance(net, torch.nn.Module):
  145.         device = list(net.parameters())[0].device
  146.     acc_sum, n = 0.0, 0
  147.     with torch.no_grad():
  148.         for X, y in data_iter:
  149.             net.eval()  # 评估模式, 关闭dropout正则化层
  150.             acc_sum += (net(X.to(device)).argmax(dim=1) == y.to(device)).float().sum().cpu().item()
  151.             net.train()  # 改回训练模式
  152.             n += y.shape[0]
  153.     return acc_sum / n