model.py

  1. import torch.nn as nn
  2. import torch
  4. class BasicBlock(nn.Module):
  5.     expansion = 1  #主分支中卷积核的个数是否发生变化
  7.     def __init__(self, in_channel, out_channel, stride = 1, downsample = None):
  8.         super(BasicBlock, self).__init__()
  9.         self.conv1 = nn.Conv2d(in_channels=in_channel, out_channels=out_channel, kernel_size=3, 
  10.                                 stride=stride, padding=1, bias=False)    # stride = 1 对应实线残差结构
  11.         # output = (input - kernel_size + 2 * padding)/stride + 1
  12.         self.bn1 = nn.BatchNorm2d(out_channel)
  13.         self.relu = nn.ReLU()
  15.         self.conv2 = nn.Conv2d(in_channels=in_channel, out_channels=out_channel, kernel_size=3, 
  16.                                 stride=1, padding=1, bias=False)
  17.         self.bn2 = nn.BatchNorm2d(out_channel)
  18.         self.downsample = downsample
  20.     def forward(self, x):    #x - 特征矩阵
  21.         identity = x    #捷径分支上的输出值
  22.         if self.downsample is not None:
  23.             identity = self.downsample(x)
  25.         out = self.conv1(x)
  26.         out = self.bn1(out)
  27.         out = self.relu(out)
  29.         out = self.conv2(out)
  30.         out = self.bn2(out)
  32.         out += identity
  33.         out = self.relu(out)
  35.         return out    # 18/34 定义完成
  38. class Bottleneck(nn.Module):
  39.     expansion = 4
  41.     def __init__(self, in_channel, out_channel, stride = 1, downsample = None):
  42.         super(Bottleneck, self).__init__()
  43.         self.conv1 = nn.Conv2d(in_channels=in_channel, out_channels=out_channel, kernel_size=1,
  44.                                 stride=1, bias=False)
  45.         self.bn1 = nn.BatchNorm2d(out_channel)
  47.         self.conv2 = nn.Conv2d(in_channels=in_channel, out_channels=out_channel, kernel_size=3,
  48.                                 stride=stride, bias=False, padding=1)    #两种不同的残差结构的步长不一致，采用传入的参数
  49.         self.bn2 = nn.BatchNorm2d(out_channel)    #padding大小一般设置为核大小的一半
  51.         self.conv3 = nn.Conv2d(in_channels=in_channel, out_channels=out_channel * self.expansion, 
  52.                                 kernel_size=1, stride=1, bias=False)
  53.         self.bn3 = nn.BatchNorm2d(out_channel * self.expansion)
  55.         self.relu = nn.ReLU(inplace=True)    #inplace：是否进行覆盖运算
  56.         self.downsample = downsample
  59.     def forward(self, x):
  60.         identity = x
  61.         if self.downsample is not None:
  62.             identity = self.downsample(x)
  64.         out = self.conv1(x)
  65.         out = self.bn1(out)
  66.         out = self.relu(out)
  68.         out = self.conv2(out)
  69.         out = self.bn2(out)
  70.         out = self.relu(out)
  72.         out = self.conv3(out)
  73.         out = self.bn3(out)
  75.         out += identity
  76.         out = self.relu(out)
  78.         return out
  80. class ResNet(nn.Module):
  81.     def __init__(self, block, blocks_num, num_classes=1000, include_top=True):    
  82.         """
  83.         block:对应残差结构，根据定义的层结构，传入不同的block
  84.         blocks_num:对应所使用残差结构的数目，是一个列表参数，e.g.34layers，3，4，6，3
  85.         include_top:便于搭建更复杂的网络
  86.         """
  87.         super(ResNet, self).__init__()
  88.         self.include_top = include_top
  89.         in_channel = 64    #通过max_pooling后得到的特征矩阵的深度
  91.         #定义第一层卷积层 
  92.         self.conv1 = nn.Conv2d(3, self.in_channel, kernel_size=7, stride=2, 
  93.                                 padding=3, bias=False)    #out_channel:卷积核个数
  94.         self.bn1 = nn.BatchNorm2d(self.in_channel)
  95.         self.relu = nn.ReLU(inplace=True)
  97.         #定义max pooling downsample
  98.         self.maxpooling = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
  100.         self.layer1 = self._make_layer(block, 64, blocks_num[0])    #Conv2对应的一系列残差结构
  101.         self.layer2 = self._make_layer(block, 128, blocks_num[0], stride=2)
  102.         self.layer3 = self._make_layer(block, 256, blocks_num[0], stride=2)
  103.         self.layer4 = self._make_layer(block, 512, blocks_num[0], stride=2)
  105.         if self.include_top:
  106.             self.avgpool = nn.AdaptiveAvgPool2d((1, 1))  # output size = (1, 1)
  107.             self.fc = nn.Linear(512 * block.expansion, num_classes)
  109.         #初始化
  110.         for m in self.modules():
  111.             if isinstance(m, nn.Conv2d):
  112.                 nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
  116.     def _make_layer(self, block, channel, block_num, stride=1):
  117.         """
  118.         channel:残差结构中主分支上第一个卷积核的个数：64，128，256，512
  119.         """
  120.         downsample = None
  121.         if stride != 1 or self.in_channel != channel * block.expansion:
  122.             #50/101/152 conv2对应的虚线残差结构只需要调整矩阵深度
  123.             downsample = nn.Sequential(
  124.                 nn.Conv2d(self.in_channel, channel * block.expansion, kernel_size=1, stride=stride, bias=False),
  125.                 nn.BatchNorm2d(channel * block.expansion))
  127.         layers = []    #空列表
  128.         layers.append(block(self.in_channel, channel, downsample = downsample, stride = stride))
  129.         self.in_channel = channel * block.expansion
  131.         for _ in range(1, block_num):
  132.             layers.append(block(self.in_channel, channel))    #输入特征矩阵的深度，主线分支上第一个卷积层的卷积核个数
  134.         return nn.Sequential(*layers)    #将list列表转换为非关键字参数传入
  136.     def forward(self, x):
  137.         x = self.conv1(x)
  138.         x = self.bn1(x)
  139.         x = self.relu(x)
  140.         x = self.maxpool(x)
  142.         x = self.layer1(x)
  143.         x = self.layer2(x)
  144.         x = self.layer3(x)
  145.         x = self.layer4(x)
  148.         if self.include_top:
  149.             x = self.avgpool(x)
  150.             x = torch.flatten(x, 1)    
  151.             x = self.fc(x)
  153.         return x
  156. def resnet34(num_classes=1000, include_top=True):
  157.     return ResNet(BasicBlock, [3, 4, 6, 3], num_classes=num_classes, include_top=include_top)
  160. def resnet101(num_classes=1000, include_top=True):
  161.     return ResNet(Bottleneck, [3, 4, 23, 3], num_classes=num_classes, include_top=include_top)

train.py

  1. import torch
  2. import torch.nn as nn
  3. from torchvision import transforms, datasets
  4. import json
  5. import matplotlib.pyplot as plt
  6. import os
  7. import torch.optim as optim
  8. from model import resnet34, resnet101
  9. import torchvision.models.resnet
  12. device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
  13. print(device)
  15. data_transform = {
  16.     "train": transforms.Compose([transforms.RandomResizedCrop(224),
  17.                                  transforms.RandomHorizontalFlip(),
  18.                                  transforms.ToTensor(),
  19.                                  transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])]),
  20.     "val": transforms.Compose([transforms.Resize(256),
  21.                                transforms.CenterCrop(224),
  22.                                transforms.ToTensor(),
  23.                                transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])}
  26. data_root = os.path.abspath(os.path.join(os.getcwd(), "../.."))  # get data root path
  27. image_path = data_root + "/data_set/flower_data/"  # flower data set path
  29. train_dataset = datasets.ImageFolder(root=image_path+"train",
  30.                                      transform=data_transform["train"])
  31. train_num = len(train_dataset)
  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)
  41. batch_size = 16
  42. train_loader = torch.utils.data.DataLoader(train_dataset,
  43.                                            batch_size=batch_size, shuffle=True,
  44.                                            num_workers=0)
  46. validate_dataset = datasets.ImageFolder(root=image_path + "val",
  47.                                         transform=data_transform["val"])
  48. val_num = len(validate_dataset)
  49. validate_loader = torch.utils.data.DataLoader(validate_dataset,
  50.                                               batch_size=batch_size, shuffle=False,
  51.                                               num_workers=0)
  53. net = resnet34()
  54. # load pretrain weights
  55. model_weight_path = "./resnet34-pre.pth"
  56. missing_keys, unexpected_keys = net.load_state_dict(torch.load(model_weight_path), strict=False)
  57. # for param in net.parameters():
  58. #     param.requires_grad = False
  59. # change fc layer structure
  60. inchannel = net.fc.in_features
  61. net.fc = nn.Linear(inchannel, 5)
  62. net.to(device)
  64. loss_function = nn.CrossEntropyLoss()
  65. optimizer = optim.Adam(net.parameters(), lr=0.0001)
  67. best_acc = 0.0
  68. save_path = './resNet34.pth'
  69. for epoch in range(3):
  70.     # train
  71.     net.train()
  72.     running_loss = 0.0
  73.     for step, data in enumerate(train_loader, start=0):
  74.         images, labels = data
  75.         optimizer.zero_grad()
  76.         logits = net(images.to(device))
  77.         loss = loss_function(logits, labels.to(device))
  78.         loss.backward()
  79.         optimizer.step()
  81.         # print statistics
  82.         running_loss += loss.item()
  83.         # print train process
  84.         rate = (step+1)/len(train_loader)
  85.         a = "*" * int(rate * 50)
  86.         b = "." * int((1 - rate) * 50)
  87.         print("\rtrain loss: {:^3.0f}%[{}->{}]{:.4f}".format(int(rate*100), a, b, loss), end="")
  88.     print()
  90.     # validate
  91.     net.eval()
  92.     acc = 0.0  # accumulate accurate number / epoch
  93.     with torch.no_grad():
  94.         for val_data in validate_loader:
  95.             val_images, val_labels = val_data
  96.             outputs = net(val_images.to(device))  # eval model only have last output layer
  97.             # loss = loss_function(outputs, test_labels)
  98.             predict_y = torch.max(outputs, dim=1)[1]
  99.             acc += (predict_y == val_labels.to(device)).sum().item()
  100.         val_accurate = acc / val_num
  101.         if val_accurate > best_acc:
  102.             best_acc = val_accurate
  103.             torch.save(net.state_dict(), save_path)
  104.         print('[epoch %d] train_loss: %.3f  test_accuracy: %.3f' %
  105.               (epoch + 1, running_loss / step, val_accurate))
  107. print('Finished Training')

predict.py

  1. import torch
  2. from model import resnet34
  3. from PIL import Image
  4. from torchvision import transforms
  5. import matplotlib.pyplot as plt
  6. import json
  8. device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
  10. data_transform = transforms.Compose(
  11.     [transforms.Resize(256),
  12.      transforms.CenterCrop(224),
  13.      transforms.ToTensor(),
  14.      transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])
  16. # load image
  17. img = Image.open("../tulip.jpg")
  18. plt.imshow(img)
  19. # [N, C, H, W]
  20. img = data_transform(img)
  21. # expand batch dimension
  22. img = torch.unsqueeze(img, dim=0)
  24. # read class_indict
  25. try:
  26.     json_file = open('./class_indices.json', 'r')
  27.     class_indict = json.load(json_file)
  28. except Exception as e:
  29.     print(e)
  30.     exit(-1)
  32. # create model
  33. model = resnet34(num_classes=5)
  34. # load model weights
  35. model_weight_path = "./resNet34.pth"
  36. model.load_state_dict(torch.load(model_weight_path, map_location=device))
  37. model.eval()
  38. with torch.no_grad():
  39.     # predict class
  40.     output = torch.squeeze(model(img))
  41.     predict = torch.softmax(output, dim=0)
  42.     predict_cla = torch.argmax(predict).numpy()
  43. print(class_indict[str(predict_cla)], predict[predict_cla].numpy())
  44. plt.show()