pytorch训练AlexNet

 1 import torch.nn as nn
 2 import torch
 3 
 4 
 5 class AlexNet(nn.Module):
 6     def __init__(self, num_classes=1000, init_weights=False):
 7         super(AlexNet, self).__init__()
 8         self.features = nn.Sequential(
 9             nn.Conv2d(3, 48, kernel_size=11, stride=4, padding=2),  # input[3, 224, 224]  output[48, 55, 55]
10             nn.ReLU(inplace=True),
11             nn.MaxPool2d(kernel_size=3, stride=2),                  # output[48, 27, 27]
12             nn.Conv2d(48, 128, kernel_size=5, padding=2),           # output[128, 27, 27]
13             nn.ReLU(inplace=True),
14             nn.MaxPool2d(kernel_size=3, stride=2),                  # output[128, 13, 13]
15             nn.Conv2d(128, 192, kernel_size=3, padding=1),          # output[192, 13, 13]
16             nn.ReLU(inplace=True),
17             nn.Conv2d(192, 192, kernel_size=3, padding=1),          # output[192, 13, 13]
18             nn.ReLU(inplace=True),
19             nn.Conv2d(192, 128, kernel_size=3, padding=1),          # output[128, 13, 13]
20             nn.ReLU(inplace=True),
21             nn.MaxPool2d(kernel_size=3, stride=2),                  # output[128, 6, 6]
22         )
23         self.classifier = nn.Sequential(
24             nn.Dropout(p=0.5),
25             nn.Linear(128 * 6 * 6, 2048),
26             nn.ReLU(inplace=True),
27             nn.Dropout(p=0.5),
28             nn.Linear(2048, 2048),
29             nn.ReLU(inplace=True),
30             nn.Linear(2048, num_classes),
31         )
32         if init_weights:
33             self._initialize_weights()
34 
35     def forward(self, x):
36         x = self.features(x)
37         x = torch.flatten(x, start_dim=1)
38         x = self.classifier(x)
39         return x
40 
41     def _initialize_weights(self):
42         for m in self.modules():
43             if isinstance(m, nn.Conv2d):
44                 nn.init.kaiming_normal_(m.weight, mode=‘fan_out‘, nonlinearity=‘relu‘)
45                 if m.bias is not None:
46                     nn.init.constant_(m.bias, 0)
47             elif isinstance(m, nn.Linear):
48                 nn.init.normal_(m.weight, 0, 0.01)
49                 nn.init.constant_(m.bias, 0)

  1 import torch
  2 import torch.nn as nn
  3 from torchvision import transforms, datasets, utils
  4 import matplotlib.pyplot as plt
  5 import numpy as np
  6 import torch.optim as optim
  7 from model import AlexNet
  8 import os
  9 import json
 10 import time
 11 
 12 
 13 def main():
 14     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 15     print("using {} device.".format(device))
 16 
 17     data_transform = {
 18         "train": transforms.Compose([transforms.RandomResizedCrop(224),
 19                                      transforms.RandomHorizontalFlip(),
 20                                      transforms.ToTensor(),
 21                                      transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]),
 22         "val": transforms.Compose([transforms.Resize((224, 224)),  # cannot 224, must (224, 224)
 23                                    transforms.ToTensor(),
 24                                    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])}
 25 
 26     data_root = os.path.abspath(os.path.join(os.getcwd(), "../.."))  # get data root path
 27     image_path = os.path.join(data_root, "data_set", "flower_data")  # flower data set path
 28     assert os.path.exists(image_path), "{} path does not exist.".format(image_path)
 29     train_dataset = datasets.ImageFolder(root=os.path.join(image_path, "train"),
 30                                          transform=data_transform["train"])
 31     train_num = len(train_dataset)
 32 
 33     # {‘daisy‘:0, ‘dandelion‘:1, ‘roses‘:2, ‘sunflower‘:3, ‘tulips‘:4}
 34     flower_list = train_dataset.class_to_idx
 35     cla_dict = dict((val, key) for key, val in flower_list.items())
 36     # write dict into json file
 37     json_str = json.dumps(cla_dict, indent=4)
 38     with open(‘class_indices.json‘, ‘w‘) as json_file:
 39         json_file.write(json_str)
 40 
 41     batch_size = 32
 42     nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0, 8])  # number of workers
 43     print(‘Using {} dataloader workers every process‘.format(nw))
 44 
 45     train_loader = torch.utils.data.DataLoader(train_dataset,
 46                                                batch_size=batch_size, shuffle=True,
 47                                                num_workers=0)
 48 
 49     validate_dataset = datasets.ImageFolder(root=os.path.join(image_path, "val"),
 50                                             transform=data_transform["val"])
 51     val_num = len(validate_dataset)
 52     validate_loader = torch.utils.data.DataLoader(validate_dataset,
 53                                                   batch_size=batch_size, shuffle=True,
 54                                                   num_workers=0)
 55 
 56     print("using {} images for training, {} images fot validation.".format(train_num,
 57                                                                            val_num))
 58     # test_data_iter = iter(validate_loader)
 59     # test_image, test_label = test_data_iter.next()
 60     # #
 61     # def imshow(img):
 62     #     img = img / 2 + 0.5  # unnormalize
 63     #     npimg = img.numpy()
 64     #     plt.imshow(np.transpose(npimg, (1, 2, 0)))
 65     #     plt.show()
 66     #
 67     # print(‘ ‘.join(‘%5s‘ % cla_dict[test_label[j].item()] for j in range(4)))
 68     # imshow(utils.make_grid(test_image))
 69 
 70     net = AlexNet(num_classes=5, init_weights=True)
 71 
 72     net.to(device)
 73     loss_function = nn.CrossEntropyLoss()
 74     # pata = list(net.parameters())
 75     optimizer = optim.Adam(net.parameters(), lr=0.0002)
 76 
 77     save_path = ‘./AlexNet.pth‘
 78     best_acc = 0.0
 79     for epoch in range(10):
 80         # train
 81         net.train()
 82         running_loss = 0.0
 83         t1 = time.perf_counter()
 84         for step, data in enumerate(train_loader, start=0):
 85             images, labels = data
 86             optimizer.zero_grad()
 87             outputs = net(images.to(device))
 88             loss = loss_function(outputs, labels.to(device))
 89             loss.backward()
 90             optimizer.step()
 91 
 92             # print statistics
 93             running_loss += loss.item()
 94             # print train process
 95             rate = (step + 1) / len(train_loader)
 96             a = "*" * int(rate * 50)
 97             b = "." * int((1 - rate) * 50)
 98             print("\rtrain loss: {:^3.0f}%[{}->{}]{:.3f}".format(int(rate * 100), a, b, loss), end="")
 99         print()
100         print(time.perf_counter()-t1)
101 
102         # validate
103         net.eval()
104         acc = 0.0  # accumulate accurate number / epoch
105         with torch.no_grad():
106             for val_data in validate_loader:
107                 val_images, val_labels = val_data
108                 outputs = net(val_images.to(device))
109                 predict_y = torch.max(outputs, dim=1)[1]
110                 acc += (predict_y == val_labels.to(device)).sum().item()
111             val_accurate = acc / val_num
112             if val_accurate > best_acc:
113                 best_acc = val_accurate
114                 torch.save(net.state_dict(), save_path)
115             print(‘[epoch %d] train_loss: %.3f  test_accuracy: %.3f‘ %
116                   (epoch + 1, running_loss / step, val_accurate))
117 
118     print(‘Finished Training‘)
119 
120 
121 if __name__ == ‘__main__‘:
122     main()