42.pdf

    1. import  torch
    2. from    torch import  nn
    3. from    torch.nn import functional as F
    4. from    torch.utils.data import DataLoader
    5. from    torchvision import datasets
    6. from    torchvision import transforms
    7. from    torch import nn, optim
    9. # from    torchvision.models import resnet18
    11. class ResBlk(nn.Module):
    12.     """
    13.     resnet block
    14.     """
    16.     def __init__(self, ch_in, ch_out):
    17.         """
    18.         :param ch_in:
    19.         :param ch_out:
    20.         """
    21.         super(ResBlk, self).__init__()
    23.         self.conv1 = nn.Conv2d(ch_in, ch_out, kernel_size=3, stride=1, padding=1)
    24.         self.bn1 = nn.BatchNorm2d(ch_out)
    25.         self.conv2 = nn.Conv2d(ch_out, ch_out, kernel_size=3, stride=1, padding=1)
    26.         self.bn2 = nn.BatchNorm2d(ch_out)
    28.         self.extra = nn.Sequential()
    29.         if ch_out != ch_in:
    30.             # [b, ch_in, h, w] => [b, ch_out, h, w]
    31.             self.extra = nn.Sequential(
    32.                 nn.Conv2d(ch_in, ch_out, kernel_size=1, stride=1),
    33.                 nn.BatchNorm2d(ch_out)
    34.             )
    37.     def forward(self, x):
    38.         """
    39.         :param x: [b, ch, h, w]
    40.         :return:
    41.         """
    42.         out = F.relu(self.bn1(self.conv1(x)))
    43.         out = self.bn2(self.conv2(out))
    44.         # short cut.
    45.         # extra module: [b, ch_in, h, w] => [b, ch_out, h, w]
    46.         # element-wise add:
    47.         out = self.extra(x) + 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, 16, kernel_size=3, stride=1, padding=1),
    61.             nn.BatchNorm2d(16)
    62.         )
    63.         # followed 4 blocks
    64.         # [b, 64, h, w] => [b, 128, h ,w]
    65.         self.blk1 = ResBlk(16, 16)
    66.         # [b, 128, h, w] => [b, 256, h, w]
    67.         self.blk2 = ResBlk(16, 32)
    68.         # # [b, 256, h, w] => [b, 512, h, w]
    69.         # self.blk3 = ResBlk(128, 256)
    70.         # # [b, 512, h, w] => [b, 1024, h, w]
    71.         # self.blk4 = ResBlk(256, 512)
    73.         self.outlayer = nn.Linear(32*32*32, 10)
    75.     def forward(self, x):
    76.         """
    77.         :param x:
    78.         :return:
    79.         """
    80.         x = F.relu(self.conv1(x))
    82.         # [b, 64, h, w] => [b, 1024, h, w]
    83.         x = self.blk1(x)
    84.         x = self.blk2(x)
    85.         # x = self.blk3(x)
    86.         # x = self.blk4(x)
    88.         # print(x.shape)
    89.         x = x.view(x.size(0), -1)
    90.         x = self.outlayer(x)
    93.         return x
    97. def main():
    98.     batchsz = 32
    100.     cifar_train = datasets.CIFAR10('cifar', True, transform=transforms.Compose([
    101.         transforms.Resize((32, 32)),
    102.         transforms.ToTensor()
    103.     ]), download=True)
    104.     cifar_train = DataLoader(cifar_train, batch_size=batchsz, shuffle=True)
    106.     cifar_test = datasets.CIFAR10('cifar', False, transform=transforms.Compose([
    107.         transforms.Resize((32, 32)),
    108.         transforms.ToTensor()
    109.     ]), download=True)
    110.     cifar_test = DataLoader(cifar_test, batch_size=batchsz, shuffle=True)
    113.     x, label = iter(cifar_train).next()
    114.     print('x:', x.shape, 'label:', label.shape)
    116.     device = torch.device('cuda')
    117.     # model = Lenet5().to(device)
    118.     model = ResNet18().to(device)
    120.     criteon = nn.CrossEntropyLoss().to(device)
    121.     optimizer = optim.Adam(model.parameters(), lr=1e-3)
    122.     print(model)
    124.     for epoch in range(1000):
    126.         model.train()
    127.         for batchidx, (x, label) in enumerate(cifar_train):
    128.             # [b, 3, 32, 32]
    129.             # [b]
    130.             x, label = x.to(device), label.to(device)
    132.             logits = model(x)
    133.             # logits: [b, 10]
    134.             # label:  [b]
    135.             # loss: tensor scalar
    136.             loss = criteon(logits, label)
    138.             # backprop
    139.             optimizer.zero_grad()
    140.             loss.backward()
    141.             optimizer.step()
    144.         #
    145.         print(epoch, 'loss:', loss.item())
    148.         model.eval()
    149.         with torch.no_grad():
    150.             # test
    151.             total_correct = 0
    152.             total_num = 0
    153.             for x, label in cifar_test:
    154.                 # [b, 3, 32, 32]
    155.                 # [b]
    156.                 x, label = x.to(device), label.to(device)
    158.                 # [b, 10]
    159.                 logits = model(x)
    160.                 # [b]
    161.                 pred = logits.argmax(dim=1)
    162.                 # [b] vs [b] => scalar tensor
    163.                 correct = torch.eq(pred, label).float().sum().item()
    164.                 total_correct += correct
    165.                 total_num += x.size(0)
    166.                 # print(correct)
    168.             acc = total_correct / total_num
    169.             print(epoch, 'acc:', acc)
    172. if __name__ == '__main__':
    173.     main()