1. import  torch
    2. from    torch import  nn
    3. from    torch.nn import functional as F
    7. class ResBlk(nn.Module):
    8.     """
    9.     resnet block
    10.     """
    12.     def __init__(self, ch_in, ch_out, stride=1):
    13.         """
    15.         :param ch_in:
    16.         :param ch_out:
    17.         """
    18.         super(ResBlk, self).__init__()
    20.         # we add stride support for resbok, which is distinct from tutorials.
    21.         self.conv1 = nn.Conv2d(ch_in, ch_out, kernel_size=3, stride=stride, padding=1)
    22.         self.bn1 = nn.BatchNorm2d(ch_out)
    23.         self.conv2 = nn.Conv2d(ch_out, ch_out, kernel_size=3, stride=1, padding=1)
    24.         self.bn2 = nn.BatchNorm2d(ch_out)
    26.         self.extra = nn.Sequential()
    27.         if ch_out != ch_in:
    28.             # [b, ch_in, h, w] => [b, ch_out, h, w]
    29.             self.extra = nn.Sequential(
    30.                 nn.Conv2d(ch_in, ch_out, kernel_size=1, stride=stride),
    31.                 nn.BatchNorm2d(ch_out)
    32.             )
    35.     def forward(self, x):
    36.         """
    38.         :param x: [b, ch, h, w]
    39.         :return:
    40.         """
    41.         out = F.relu(self.bn1(self.conv1(x)))
    42.         out = self.bn2(self.conv2(out))
    43.         # short cut.
    44.         # extra module: [b, ch_in, h, w] => [b, ch_out, h, w]
    45.         # element-wise add:
    46.         out = self.extra(x) + out
    47.         out = F.relu(out)
    49.         return out
    54. class ResNet18(nn.Module):
    56.     def __init__(self):
    57.         super(ResNet18, self).__init__()
    59.         self.conv1 = nn.Sequential(
    60.             nn.Conv2d(3, 64, kernel_size=3, stride=3, padding=0),
    61.             nn.BatchNorm2d(64)
    62.         )
    63.         # followed 4 blocks
    64.         # [b, 64, h, w] => [b, 128, h ,w]
    65.         self.blk1 = ResBlk(64, 128, stride=2)
    66.         # [b, 128, h, w] => [b, 256, h, w]
    67.         self.blk2 = ResBlk(128, 256, stride=2)
    68.         # # [b, 256, h, w] => [b, 512, h, w]
    69.         self.blk3 = ResBlk(256, 512, stride=2)
    70.         # # [b, 512, h, w] => [b, 1024, h, w]
    71.         self.blk4 = ResBlk(512, 512, stride=2)
    73.         self.outlayer = nn.Linear(512*1*1, 10)
    75.     def forward(self, x):
    76.         """
    78.         :param x:
    79.         :return:
    80.         """
    81.         x = F.relu(self.conv1(x))
    83.         # [b, 64, h, w] => [b, 1024, h, w]
    84.         x = self.blk1(x)
    85.         x = self.blk2(x)
    86.         x = self.blk3(x)
    87.         x = self.blk4(x)
    90.         # print('after conv:', x.shape) #[b, 512, 2, 2]
    91.         # [b, 512, h, w] => [b, 512, 1, 1]
    92.         x = F.adaptive_avg_pool2d(x, [1, 1])
    93.         # print('after pool:', x.shape)
    94.         x = x.view(x.size(0), -1)
    95.         x = self.outlayer(x)
    98.         return x
    102. def main():
    104.     blk = ResBlk(64, 128, stride=4)
    105.     tmp = torch.randn(2, 64, 32, 32)
    106.     out = blk(tmp)
    107.     print('block:', out.shape)
    109.     x = torch.randn(2, 3, 32, 32)
    110.     model = ResNet18()
    111.     out = model(x)
    112.     print('resnet:', out.shape)
    117. if __name__ == '__main__':
    118.     main()
    1. import  torch
    2. from    torch.utils.data import DataLoader
    3. from    torchvision import datasets
    4. from    torchvision import transforms
    5. from    torch import nn, optim
    7. from    lenet5 import Lenet5
    8. from    resnet import ResNet18
    10. def main():
    11.     batchsz = 128
    13.     cifar_train = datasets.CIFAR10('cifar', True, transform=transforms.Compose([
    14.         transforms.Resize((32, 32)),
    15.         transforms.ToTensor(),
    16.         transforms.Normalize(mean=[0.485, 0.456, 0.406],
    17.                              std=[0.229, 0.224, 0.225])
    18.     ]), download=True)
    19.     cifar_train = DataLoader(cifar_train, batch_size=batchsz, shuffle=True)
    21.     cifar_test = datasets.CIFAR10('cifar', False, transform=transforms.Compose([
    22.         transforms.Resize((32, 32)),
    23.         transforms.ToTensor(),
    24.         transforms.Normalize(mean=[0.485, 0.456, 0.406],
    25.                              std=[0.229, 0.224, 0.225])
    26.     ]), download=True)
    27.     cifar_test = DataLoader(cifar_test, batch_size=batchsz, shuffle=True)
    30.     x, label = iter(cifar_train).next()
    31.     print('x:', x.shape, 'label:', label.shape)
    33.     device = torch.device('cuda')
    34.     # model = Lenet5().to(device)
    35.     model = ResNet18().to(device)
    37.     criteon = nn.CrossEntropyLoss().to(device)
    38.     optimizer = optim.Adam(model.parameters(), lr=1e-3)
    39.     print(model)
    41.     for epoch in range(1000):
    43.         model.train()
    44.         for batchidx, (x, label) in enumerate(cifar_train):
    45.             # [b, 3, 32, 32]
    46.             # [b]
    47.             x, label = x.to(device), label.to(device)
    50.             logits = model(x)
    51.             # logits: [b, 10]
    52.             # label:  [b]
    53.             # loss: tensor scalar
    54.             loss = criteon(logits, label)
    56.             # backprop
    57.             optimizer.zero_grad()
    58.             loss.backward()
    59.             optimizer.step()
    62.         print(epoch, 'loss:', loss.item())
    65.         model.eval()
    66.         with torch.no_grad():
    67.             # test
    68.             total_correct = 0
    69.             total_num = 0
    70.             for x, label in cifar_test:
    71.                 # [b, 3, 32, 32]
    72.                 # [b]
    73.                 x, label = x.to(device), label.to(device)
    75.                 # [b, 10]
    76.                 logits = model(x)
    77.                 # [b]
    78.                 pred = logits.argmax(dim=1)
    79.                 # [b] vs [b] => scalar tensor
    80.                 correct = torch.eq(pred, label).float().sum().item()
    81.                 total_correct += correct
    82.                 total_num += x.size(0)
    83.                 # print(correct)
    85.             acc = total_correct / total_num
    86.             print(epoch, 'test acc:', acc)
    90. if __name__ == '__main__':
    91.     main()
    1. import  torch
    2. from    torch import nn
    3. from    torch.nn import functional as F
    8. class Lenet5(nn.Module):
    9.     """
    10.     for cifar10 dataset.
    11.     """
    12.     def __init__(self):
    13.         super(Lenet5, self).__init__()
    15.         self.conv_unit = nn.Sequential(
    16.             # x: [b, 3, 32, 32] => [b, 16, ]
    17.             nn.Conv2d(3, 16, kernel_size=5, stride=1, padding=0),
    18.             nn.MaxPool2d(kernel_size=2, stride=2, padding=0),
    19.             #
    20.             nn.Conv2d(16, 32, kernel_size=5, stride=1, padding=0),
    21.             nn.MaxPool2d(kernel_size=2, stride=2, padding=0),
    22.             #
    23.         )
    24.         # flatten
    25.         # fc unit
    26.         self.fc_unit = nn.Sequential(
    27.             nn.Linear(32*5*5, 32),
    28.             nn.ReLU(),
    29.             # nn.Linear(120, 84),
    30.             # nn.ReLU(),
    31.             nn.Linear(32, 10)
    32.         )
    35.         # [b, 3, 32, 32]
    36.         tmp = torch.randn(2, 3, 32, 32)
    37.         out = self.conv_unit(tmp)
    38.         # [b, 16, 5, 5]
    39.         print('conv out:', out.shape)
    41.         # # use Cross Entropy Loss
    42.         # self.criteon = nn.CrossEntropyLoss()
    46.     def forward(self, x):
    47.         """
    49.         :param x: [b, 3, 32, 32]
    50.         :return:
    51.         """
    52.         batchsz = x.size(0)
    53.         # [b, 3, 32, 32] => [b, 16, 5, 5]
    54.         x = self.conv_unit(x)
    55.         # [b, 16, 5, 5] => [b, 16*5*5]
    56.         x = x.view(batchsz, 32*5*5)
    57.         # [b, 16*5*5] => [b, 10]
    58.         logits = self.fc_unit(x)
    60.         # # [b, 10]
    61.         # pred = F.softmax(logits, dim=1)
    62.         # loss = self.criteon(logits, y)
    64.         return logits
    66. def main():
    68.     net = Lenet5()
    70.     tmp = torch.randn(2, 3, 32, 32)
    71.     out = net(tmp)
    72.     print('lenet out:', out.shape)
    77. if __name__ == '__main__':
    78.     main()