【PyTorch学习笔记】9.自定义数据集

男娘i 2022-09-05 13:58 250阅读 0赞

### 文章目录 ###

*  59.自定义数据集
 *  60.数据预处理
 *  61.迁移学习
 *  62.实战
 *   *  62.1RuntimeError解决
     *  62.2迁移学习实战

根据龙良曲Pytorch学习视频整理，视频链接：  
[【计算机-AI】PyTorch学这个就够了！][-AI_PyTorch]  
[(好课推荐)深度学习与PyTorch入门实战——主讲人龙良曲][PyTorch]

# 59.自定义数据集 #

4 steps

1.  Load data  
    Inherit from torch.utils.data.Dataset  
    \_\_len\_\_  
    \_\_getitem\_\_
2.  Build model  
    Inherit from base class  
    Define forward graph
3.  Train and Test
4.  Tranfer Learning

# 60.数据预处理 #

Preprocessing

*  image Resize  
    224X224 for ResNet18
 *  Data Argumentation  
    Rotate, Crop
 *  Normalize  
    Mean, std
 *  ToTensor

# 61.迁移学习 #

![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2tlaXZlbl8_size_16_color_FFFFFF_t_70]  
![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2tlaXZlbl8_size_16_color_FFFFFF_t_70 1]

# 62.实战 #

## 62.1RuntimeError解决 ##

这里如果我使用多卡加载数据会报错，实战时写了下面一行代码：

train_loader = DataLoader(train_db, batch_size=batchsz, shuffle=True, num_workers=4)

其中`num_workers=4` 表示cuda多卡并行运行，运行时报错：

RuntimeError: cuDNN error: CUDNN_STATUS_INTERNAL_ERROR

查阅资料后解决：需要指定目标显卡参数：

os.environ['CUDA_VISIBLE_DEVICES'] = '1, 2, 3, 4'

> 参考博客：  
> [PyTorch指定GPU的方法][PyTorch_GPU]  
> [pytorch 使用指定的GPU RuntimeError: CUDA error: invalid device ordinal][pytorch _GPU RuntimeError_ CUDA error_ invalid device ordinal]

## 62.2迁移学习实战 ##

没有使用迁移学习时：

best acc: 0.8583690987124464 best epoch: 9
    test acc: 0.905982905982906

使用迁移学习：

best acc: 0.9098712446351931 best epoch: 5
    test acc: 0.9358974358974359

![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2tlaXZlbl8_size_16_color_FFFFFF_t_70 2]

**pokemon.py**

import torch
    import os, glob
    import random, csv
    from torch.utils.data import Dataset, DataLoader
    import torchvision
    from torchvision import transforms
    from PIL import Image
    import visdom
    import time
    
    class Pokemon(Dataset):
        def __init__(self, root, resize, mode):
            super(Pokemon, self).__init__()
    
            self.root = root
            self.resize = resize
    
            self.name2label = { }
            for name in sorted(os.listdir(os.path.join(root))):
                if not os.path.isdir(os.path.join(root, name)):
                    continue    # 如果不是文件夹就跳过
    
                self.name2label[name] = len(self.name2label.keys())
    
            # print(self.name2label) # {'bulbasaur': 0, 'charmander': 1, 'mewtwo': 2, 'pikachu': 3, 'squirtle': 4}
    
            # image, label
            self.images, self.labels = self.load_csv('images.csv')
    
            if mode == 'train':     # 60%
                self.images = self.images[:int(0.6*len(self.images))]
                self.labels = self.labels[:int(0.6*len(self.labels))]
            elif mode == 'val':  # 20%
                self.images = self.images[int(0.6 * len(self.images)):int(0.8 * len(self.images))]
                self.labels = self.labels[int(0.6 * len(self.labels)):int(0.8 * len(self.labels))]
            elif mode == 'test':  # 20%
                self.images = self.images[int(0.8 * len(self.images)):]
                self.labels = self.labels[int(0.8 * len(self.labels)):]
    
        def load_csv(self, filename):
            if not os.path.exists(os.path.join(self.root, filename)):
                images = []
                for name in self.name2label.keys():
                    # 'pokemon/bulbasaur/00000000.png'
                    images += glob.glob(os.path.join(self.root, name, '*.png'))
                    images += glob.glob(os.path.join(self.root, name, '*.jpg'))
                    images += glob.glob(os.path.join(self.root, name, '*.jpeg'))
    
                print(len(images), images)  # 1167
    
                random.shuffle(images)
                with open(os.path.join(self.root, filename), mode='w', newline='') as f:
                    writer = csv.writer(f)
                    for img in images:
                        name = img.split(os.sep)[-2]
                        label = self.name2label[name]
                        writer.writerow([img, label])
                    print('writen into csv file:', filename)
    
            # read from csv file
            images, labels = [], []
            with open(os.path.join(self.root, filename)) as f:
                reader = csv.reader(f)
                for row in reader:
                    img, label = row
                    label = int(label)
    
                    images.append(img)
                    labels.append(label)
    
            assert len(images) == len(labels)
    
            return images, labels
    
    
        def __len__(self):
            return len(self.images)
    
    
        def denormalize(self, x_hat):
            mean = [0.485, 0.456, 0.406]
            std = [0.229, 0.224, 0.225]
    
            # x_hat = (x - mean) / std
            # x = x_hat * std + mean
            # x: [c, h, w]
            # mean: [3] => [3, 1, 1]
            mean = torch.tensor(mean).unsqueeze(1).unsqueeze(1)
            std = torch.tensor(std).unsqueeze(1).unsqueeze(1)
    
            x = x_hat * std + mean
    
            return x
    
    
        def __getitem__(self, idx):
            # idx~[0, len(images)]
            img, label = self.images[idx], self.labels[idx]
    
            tf = transforms.Compose([
                lambda x: Image.open(x).convert('RGB'),  # string => image data
                transforms.Resize((int(self.resize*1.25), int(self.resize*1.25))),
                transforms.RandomRotation(15),
                transforms.CenterCrop(self.resize),
                transforms.ToTensor(),
                transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                     std=[0.229, 0.224, 0.225])
            ])
    
            img = tf(img)
            label = torch.tensor(label)
    
            return img, label
    
    def main():
        viz = visdom.Visdom()
    
        print(os.path)
        path = 'G:/Project/PYTHON/Demo/Pytorch21_7_29/custom_dataset/pokemon'
        db = Pokemon(path, 224, 'train')
    
        x, y = next(iter(db))
        print('sample:', x.shape, y.shape, y)
    
        viz.image(db.denormalize(x), win='sample_x', opts=dict(title='sample_x'))
    
        # tf = transforms.Compose([
        # transforms.Resize(64, 64),
        # transforms.ToTensor()
        # ])
        loader = DataLoader(db, batch_size=32, shuffle=True)
        for x, y in loader:
            viz.images(db.denormalize(x), nrow=8, win='batch', opts=dict(title='batch'))
            viz.text(str(y.numpy()), win='label', opts=dict(title='batch-y'))
            time.sleep(10)
    
    if __name__ == '__main__':
        main()

**utils.py**

import torch
    from torch import nn
    import matplotlib.pyplot as plt
    
    class Flatten(nn.Module):
        def __init__(self):
            super(Flatten, self).__init__()
    
        def foward(self, x):
            shape = torch.prod(torch.tensor(x.shape[1:])).item()
            return x.view(-1, shape)
    
    def plot_image(img, label, name):
        fig = plt.figure()
        for i in range(6):
            plt.subplot(2, 3, i+1)
            plt.tight_layout()
            plt.imshow(img[i][0]*0.3081+0.1307, cmap='gray', interpolation='none')
            plt.title('{}: {}'.format(name, label[i].item()))
            plt.xticks([])
            plt.yticks([])
        plt.show()

**train\_tranfer.py**

import torch
    from torch import nn, optim
    import visdom
    from torch.utils.data import DataLoader
    from Pytorch21_7_29.custom_dataset.pokemon import Pokemon
    from Pytorch21_7_29.custom_dataset.resnet import ResNet18
    from torchvision.models import resnet18
    from Pytorch21_7_29.custom_dataset.utils import Flatten
    import os
    os.environ["CUDA_VISIBLE_DEVICES"] = '0,1,2,3'
    
    batchsz = 32
    lr = 1e-3
    epochs = 10
    
    device = torch.device('cuda')
    torch.manual_seed(1234)
    
    train_db = Pokemon('pokemon', 224, mode='train')
    val_db = Pokemon('pokemon', 224, mode='val')
    test_db = Pokemon('pokemon', 224, mode='test')
    
    train_loader = DataLoader(train_db, batch_size=batchsz, shuffle=True,
                              num_workers=3)    # 并行
    val_loader = DataLoader(val_db, batch_size=batchsz, shuffle=True,
                              num_workers=2)
    test_loader = DataLoader(test_db, batch_size=batchsz, shuffle=True,
                              num_workers=2)
    
    viz = visdom.Visdom()
    
    def evalute(model, loader):
        correct = 0
        total = len(loader.dataset)
    
        for x, y in loader:
            x, y = x.to(device), y.to(device)
            with torch.no_grad():
                logits = model(x)
                pred = logits.argmax(dim=1)
            correct += torch.eq(pred, y).sum().float().item()
    
        return correct / total
    
    
    def main():
        # model = ResNet18(5).to(device)
    
        trained_model = resnet18(pretrained=True)
        model = nn.Sequential(
            *list(trained_model.children())[:-1],    # [b, 512, 1, 1]
            Flatten(),      # [b, 512, 1, 1] => [b, 512]
            nn.Linear(512, 5)
        ).to(device)
    
        optimizer = optim.Adam(model.parameters(), lr=lr)
        criteon = nn.CrossEntropyLoss()
    
        best_acc, best_epoch = 0, 0
        global_step = 0
    
        viz.line([0], [0], win='loss', opts=dict(title='loss'))
        viz.line([0], [0], win='val_acc', opts=dict(title='val_acc'))
        for epoch in range(epochs):
            for step, (x, y) in enumerate(train_loader):
                # x: [b, 3, 224, 224] y: [b]
                x, y = x.to(device), y.to(device)
    
                logits = model(x)
                loss = criteon(logits, y)
    
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()
    
                viz.line([loss.item()], [global_step], win='loss', update='append')
                global_step += 1
    
            val_acc = evalute(model, val_loader)
            if val_acc > best_acc:
                best_epoch = epoch
                best_acc = val_acc
    
                torch.save(model.state_dict(), 'best.mdl')
    
                viz.line([val_acc], [global_step], win='val_acc', update='append')
    
        print('best acc:', best_acc, 'best epoch:', best_epoch)
    
        root_path = os.getcwd()
        model.load_state_dict(torch.load(root_path + '/best.mdl'))
        print('loaded from ckpt!')
    
        test_acc = evalute(model, test_loader)
        print('test acc:', test_acc)
    
    
    if __name__ == '__main__':
        main()

[-AI_PyTorch]: https://www.bilibili.com/video/BV1Rv411y7oE?p=1
[PyTorch]: https://www.bilibili.com/video/BV1mh41167e3?p=1
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2tlaXZlbl8_size_16_color_FFFFFF_t_70]: /images/20220829/abf44e5713184b0b8ecb3b2af3b91293.png
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2tlaXZlbl8_size_16_color_FFFFFF_t_70 1]: /images/20220829/fc76a45af74040949590a0003cb5e7d3.png
[PyTorch_GPU]: https://blog.csdn.net/m0_37613352/article/details/118513705
[pytorch _GPU RuntimeError_ CUDA error_ invalid device ordinal]: https://junchu.blog.csdn.net/article/details/107720675
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2tlaXZlbl8_size_16_color_FFFFFF_t_70 2]: /images/20220829/99fb0a8616d94fab9c08eea1ea3076d4.png