宝可梦数据集.pdf

    1. import  torch
    2. import  os, glob
    3. import  random, csv
    5. from    torch.utils.data import Dataset, DataLoader
    7. from    torchvision import transforms
    8. from    PIL import Image
    11. class Pokemon(Dataset):
    13.     def __init__(self, root, resize, mode):
    14.         super(Pokemon, self).__init__()
    16.         self.root = root
    17.         self.resize = resize
    19.         self.name2label = {} # "sq...":0
    20.         for name in sorted(os.listdir(os.path.join(root))):
    21.             if not os.path.isdir(os.path.join(root, name)):
    22.                 continue
    24.             self.name2label[name] = len(self.name2label.keys())
    26.         # print(self.name2label)
    28.         # image, label
    29.         self.images, self.labels = self.load_csv('images.csv')
    31.         if mode=='train': # 60%
    32.             self.images = self.images[:int(0.6*len(self.images))]
    33.             self.labels = self.labels[:int(0.6*len(self.labels))]
    34.         elif mode=='val': # 20% = 60%->80%
    35.             self.images = self.images[int(0.6*len(self.images)):int(0.8*len(self.images))]
    36.             self.labels = self.labels[int(0.6*len(self.labels)):int(0.8*len(self.labels))]
    37.         else: # 20% = 80%->100%
    38.             self.images = self.images[int(0.8*len(self.images)):]
    39.             self.labels = self.labels[int(0.8*len(self.labels)):]
    45.     def load_csv(self, filename):
    47.         if not os.path.exists(os.path.join(self.root, filename)):
    48.             images = []
    49.             for name in self.name2label.keys():
    50.                 # 'pokemon\\mewtwo\\00001.png
    51.                 images += glob.glob(os.path.join(self.root, name, '*.png'))
    52.                 images += glob.glob(os.path.join(self.root, name, '*.jpg'))
    53.                 images += glob.glob(os.path.join(self.root, name, '*.jpeg'))
    55.             # 1167, 'pokemon\\bulbasaur\\00000000.png'
    56.             print(len(images), images)
    58.             random.shuffle(images)
    59.             with open(os.path.join(self.root, filename), mode='w', newline='') as f:
    60.                 writer = csv.writer(f)
    61.                 for img in images: # 'pokemon\\bulbasaur\\00000000.png'
    62.                     name = img.split(os.sep)[-2]
    63.                     label = self.name2label[name]
    64.                     # 'pokemon\\bulbasaur\\00000000.png', 0
    65.                     writer.writerow([img, label])
    66.                 print('writen into csv file:', filename)
    68.         # read from csv file
    69.         images, labels = [], []
    70.         with open(os.path.join(self.root, filename)) as f:
    71.             reader = csv.reader(f)
    72.             for row in reader:
    73.                 # 'pokemon\\bulbasaur\\00000000.png', 0
    74.                 img, label = row
    75.                 label = int(label)
    77.                 images.append(img)
    78.                 labels.append(label)
    80.         assert len(images) == len(labels)
    82.         return images, labels
    86.     def __len__(self):
    88.         return len(self.images)
    91.     def denormalize(self, x_hat):
    93.         mean = [0.485, 0.456, 0.406]
    94.         std = [0.229, 0.224, 0.225]
    96.         # x_hat = (x-mean)/std
    97.         # x = x_hat*std = mean
    98.         # x: [c, h, w]
    99.         # mean: [3] => [3, 1, 1]
    100.         mean = torch.tensor(mean).unsqueeze(1).unsqueeze(1)
    101.         std = torch.tensor(std).unsqueeze(1).unsqueeze(1)
    102.         # print(mean.shape, std.shape)
    103.         x = x_hat * std + mean
    105.         return x
    108.     def __getitem__(self, idx):
    109.         # idx~[0~len(images)]
    110.         # self.images, self.labels
    111.         # img: 'pokemon\\bulbasaur\\00000000.png'
    112.         # label: 0
    113.         img, label = self.images[idx], self.labels[idx]
    115.         tf = transforms.Compose([
    116.             lambda x:Image.open(x).convert('RGB'), # string path= > image data
    117.             transforms.Resize((int(self.resize*1.25), int(self.resize*1.25))),
    118.             transforms.RandomRotation(15),
    119.             transforms.CenterCrop(self.resize),
    120.             transforms.ToTensor(),
    121.             transforms.Normalize(mean=[0.485, 0.456, 0.406],
    122.                                  std=[0.229, 0.224, 0.225])
    123.         ])
    125.         img = tf(img)
    126.         label = torch.tensor(label)
    129.         return img, label
    135. def main():
    137.     import  visdom
    138.     import  time
    139.     import  torchvision
    141.     viz = visdom.Visdom()
    143.     # tf = transforms.Compose([
    144.     #                 transforms.Resize((64,64)),
    145.     #                 transforms.ToTensor(),
    146.     # ])
    147.     # db = torchvision.datasets.ImageFolder(root='pokemon', transform=tf)
    148.     # loader = DataLoader(db, batch_size=32, shuffle=True)
    149.     #
    150.     # print(db.class_to_idx)
    151.     #
    152.     # for x,y in loader:
    153.     #     viz.images(x, nrow=8, win='batch', opts=dict(title='batch'))
    154.     #     viz.text(str(y.numpy()), win='label', opts=dict(title='batch-y'))
    155.     #
    156.     #     time.sleep(10)
    159.     db = Pokemon('pokemon', 64, 'train')
    161.     x,y = next(iter(db))
    162.     print('sample:', x.shape, y.shape, y)
    164.     viz.image(db.denormalize(x), win='sample_x', opts=dict(title='sample_x'))
    166.     loader = DataLoader(db, batch_size=32, shuffle=True, num_workers=8)
    168.     for x,y in loader:
    169.         viz.images(db.denormalize(x), nrow=8, win='batch', opts=dict(title='batch'))
    170.         viz.text(str(y.numpy()), win='label', opts=dict(title='batch-y'))
    172.         time.sleep(10)
    174. if __name__ == '__main__':
    175.     main()
    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.         """
    14.         :param ch_in:
    15.         :param ch_out:
    16.         """
    17.         super(ResBlk, self).__init__()
    19.         self.conv1 = nn.Conv2d(ch_in, ch_out, kernel_size=3, stride=stride, padding=1)
    20.         self.bn1 = nn.BatchNorm2d(ch_out)
    21.         self.conv2 = nn.Conv2d(ch_out, ch_out, kernel_size=3, stride=1, padding=1)
    22.         self.bn2 = nn.BatchNorm2d(ch_out)
    24.         self.extra = nn.Sequential()
    25.         if ch_out != ch_in:
    26.             # [b, ch_in, h, w] => [b, ch_out, h, w]
    27.             self.extra = nn.Sequential(
    28.                 nn.Conv2d(ch_in, ch_out, kernel_size=1, stride=stride),
    29.                 nn.BatchNorm2d(ch_out)
    30.             )
    33.     def forward(self, x):
    34.         """
    35.         :param x: [b, ch, h, w]
    36.         :return:
    37.         """
    38.         out = F.relu(self.bn1(self.conv1(x)))
    39.         out = self.bn2(self.conv2(out))
    40.         # short cut.
    41.         # extra module: [b, ch_in, h, w] => [b, ch_out, h, w]
    42.         # element-wise add:
    43.         out = self.extra(x) + out
    44.         out = F.relu(out)
    46.         return out
    51. class ResNet18(nn.Module):
    53.     def __init__(self, num_class):
    54.         super(ResNet18, self).__init__()
    56.         self.conv1 = nn.Sequential(
    57.             nn.Conv2d(3, 16, kernel_size=3, stride=3, padding=0),
    58.             nn.BatchNorm2d(16)
    59.         )
    60.         # followed 4 blocks
    61.         # [b, 16, h, w] => [b, 32, h ,w]
    62.         self.blk1 = ResBlk(16, 32, stride=3)
    63.         # [b, 32, h, w] => [b, 64, h, w]
    64.         self.blk2 = ResBlk(32, 64, stride=3)
    65.         # # [b, 64, h, w] => [b, 128, h, w]
    66.         self.blk3 = ResBlk(64, 128, stride=2)
    67.         # # [b, 128, h, w] => [b, 256, h, w]
    68.         self.blk4 = ResBlk(128, 256, stride=2)
    70.         # [b, 256, 7, 7]
    71.         self.outlayer = nn.Linear(256*3*3, num_class)
    73.     def forward(self, x):
    74.         """
    75.         :param x:
    76.         :return:
    77.         """
    78.         x = F.relu(self.conv1(x))
    80.         # [b, 64, h, w] => [b, 1024, h, w]
    81.         x = self.blk1(x)
    82.         x = self.blk2(x)
    83.         x = self.blk3(x)
    84.         x = self.blk4(x)
    86.         # print(x.shape)
    87.         x = x.view(x.size(0), -1)
    88.         x = self.outlayer(x)
    91.         return x
    95. def main():
    96.     blk = ResBlk(64, 128)
    97.     tmp = torch.randn(2, 64, 224, 224)
    98.     out = blk(tmp)
    99.     print('block:', out.shape)
    102.     model = ResNet18(5)
    103.     tmp = torch.randn(2, 3, 224, 224)
    104.     out = model(tmp)
    105.     print('resnet:', out.shape)
    107.     p = sum(map(lambda p:p.numel(), model.parameters()))
    108.     print('parameters size:', p)
    111. if __name__ == '__main__':
    112.     main()
    1. import  torch
    2. from    torch import optim, nn
    3. import  visdom
    4. import  torchvision
    5. from    torch.utils.data import DataLoader
    7. from    pokemon import Pokemon
    8. from    resnet import ResNet18
    12. batchsz = 32
    13. lr = 1e-3
    14. epochs = 10
    16. device = torch.device('cuda')
    17. torch.manual_seed(1234)
    20. train_db = Pokemon('pokemon', 224, mode='train')
    21. val_db = Pokemon('pokemon', 224, mode='val')
    22. test_db = Pokemon('pokemon', 224, mode='test')
    23. train_loader = DataLoader(train_db, batch_size=batchsz, shuffle=True,
    24.                           num_workers=4)
    25. val_loader = DataLoader(val_db, batch_size=batchsz, num_workers=2)
    26. test_loader = DataLoader(test_db, batch_size=batchsz, num_workers=2)
    29. viz = visdom.Visdom()
    31. def evalute(model, loader):
    32.     model.eval()
    34.     correct = 0
    35.     total = len(loader.dataset)
    37.     for x,y in loader:
    38.         x,y = x.to(device), y.to(device)
    39.         with torch.no_grad():
    40.             logits = model(x)
    41.             pred = logits.argmax(dim=1)
    42.         correct += torch.eq(pred, y).sum().float().item()
    44.     return correct / total
    46. def main():
    48.     model = ResNet18(5).to(device)
    49.     optimizer = optim.Adam(model.parameters(), lr=lr)
    50.     criteon = nn.CrossEntropyLoss()
    53.     best_acc, best_epoch = 0, 0
    54.     global_step = 0
    55.     viz.line([0], [-1], win='loss', opts=dict(title='loss'))
    56.     viz.line([0], [-1], win='val_acc', opts=dict(title='val_acc'))
    57.     for epoch in range(epochs):
    59.         for step, (x,y) in enumerate(train_loader):
    61.             # x: [b, 3, 224, 224], y: [b]
    62.             x, y = x.to(device), y.to(device)
    64.             model.train()
    65.             logits = model(x)
    66.             loss = criteon(logits, y)
    68.             optimizer.zero_grad()
    69.             loss.backward()
    70.             optimizer.step()
    72.             viz.line([loss.item()], [global_step], win='loss', update='append')
    73.             global_step += 1
    75.         if epoch % 1 == 0:
    77.             val_acc = evalute(model, val_loader)
    78.             if val_acc> best_acc:
    79.                 best_epoch = epoch
    80.                 best_acc = val_acc
    82.                 torch.save(model.state_dict(), 'best.mdl')
    84.                 viz.line([val_acc], [global_step], win='val_acc', update='append')
    87.     print('best acc:', best_acc, 'best epoch:', best_epoch)
    89.     model.load_state_dict(torch.load('best.mdl'))
    90.     print('loaded from ckpt!')
    92.     test_acc = evalute(model, test_loader)
    93.     print('test acc:', test_acc)
    99. if __name__ == '__main__':
    100.     main()
    1. import  torch
    2. from    torch import optim, nn
    3. import  visdom
    4. import  torchvision
    5. from    torch.utils.data import DataLoader
    7. from    pokemon import Pokemon
    8. # from    resnet import ResNet18
    9. from    torchvision.models import resnet18
    11. from    utils import Flatten
    13. batchsz = 32
    14. lr = 1e-3
    15. epochs = 10
    17. device = torch.device('cuda')
    18. torch.manual_seed(1234)
    21. train_db = Pokemon('pokemon', 224, mode='train')
    22. val_db = Pokemon('pokemon', 224, mode='val')
    23. test_db = Pokemon('pokemon', 224, mode='test')
    24. train_loader = DataLoader(train_db, batch_size=batchsz, shuffle=True,
    25.                           num_workers=4)
    26. val_loader = DataLoader(val_db, batch_size=batchsz, num_workers=2)
    27. test_loader = DataLoader(test_db, batch_size=batchsz, num_workers=2)
    30. viz = visdom.Visdom()
    32. def evalute(model, loader):
    33.     model.eval()
    35.     correct = 0
    36.     total = len(loader.dataset)
    38.     for x,y in loader:
    39.         x,y = x.to(device), y.to(device)
    40.         with torch.no_grad():
    41.             logits = model(x)
    42.             pred = logits.argmax(dim=1)
    43.         correct += torch.eq(pred, y).sum().float().item()
    45.     return correct / total
    47. def main():
    49.     # model = ResNet18(5).to(device)
    50.     trained_model = resnet18(pretrained=True)
    51.     model = nn.Sequential(*list(trained_model.children())[:-1], #[b, 512, 1, 1]
    52.                           Flatten(), # [b, 512, 1, 1] => [b, 512]
    53.                           nn.Linear(512, 5)
    54.                           ).to(device)
    55.     # x = torch.randn(2, 3, 224, 224)
    56.     # print(model(x).shape)
    58.     optimizer = optim.Adam(model.parameters(), lr=lr)
    59.     criteon = nn.CrossEntropyLoss()
    62.     best_acc, best_epoch = 0, 0
    63.     global_step = 0
    64.     viz.line([0], [-1], win='loss', opts=dict(title='loss'))
    65.     viz.line([0], [-1], win='val_acc', opts=dict(title='val_acc'))
    66.     for epoch in range(epochs):
    68.         for step, (x,y) in enumerate(train_loader):
    70.             # x: [b, 3, 224, 224], y: [b]
    71.             x, y = x.to(device), y.to(device)
    73.             model.train()
    74.             logits = model(x)
    75.             loss = criteon(logits, y)
    77.             optimizer.zero_grad()
    78.             loss.backward()
    79.             optimizer.step()
    81.             viz.line([loss.item()], [global_step], win='loss', update='append')
    82.             global_step += 1
    84.         if epoch % 1 == 0:
    86.             val_acc = evalute(model, val_loader)
    87.             if val_acc> best_acc:
    88.                 best_epoch = epoch
    89.                 best_acc = val_acc
    91.                 torch.save(model.state_dict(), 'best.mdl')
    93.                 viz.line([val_acc], [global_step], win='val_acc', update='append')
    96.     print('best acc:', best_acc, 'best epoch:', best_epoch)
    98.     model.load_state_dict(torch.load('best.mdl'))
    99.     print('loaded from ckpt!')
    101.     test_acc = evalute(model, test_loader)
    102.     print('test acc:', test_acc)
    108. if __name__ == '__main__':
    109.     main()
    1. from    matplotlib import pyplot as plt
    2. import  torch
    3. from    torch import nn
    5. class Flatten(nn.Module):
    7.     def __init__(self):
    8.         super(Flatten, self).__init__()
    10.     def forward(self, x):
    11.         shape = torch.prod(torch.tensor(x.shape[1:])).item()
    12.         return x.view(-1, shape)
    15. def plot_image(img, label, name):
    17.     fig = plt.figure()
    18.     for i in range(6):
    19.         plt.subplot(2, 3, i + 1)
    20.         plt.tight_layout()
    21.         plt.imshow(img[i][0]*0.3081+0.1307, cmap='gray', interpolation='none')
    22.         plt.title("{}: {}".format(name, label[i].item()))
    23.         plt.xticks([])
    24.         plt.yticks([])
    25.     plt.show()