pytorch-day09（自定义数据集 & 迁移学习）

1、自定义数据集

"""
自定义数据集的基础操作
"""
import torch
from torch.utils.data import Dataset

class Pokemon(Dataset):  # 继承Datastet类，如自定义模型继承Module类一样
    def __init__(self):
        super(Pokemon, self).__init__()
        pass

    def __len__(self):  
        pass

    def __getitem__(self, idx): 
        pass  

　　例如：

import torch
from torch.utils.data import Dataset

class NumbersDataset(Dataset):  # 继承Datastet类，如自定义模型继承Module类一样
    def __init__(self, training=True):
        if training:
            self.samples = list(range(1, 1001))  # 训练数据集
        else:
            self.samples = list(range(1001, 1501))

    def __len__(self):  # 返回元素的个数(数据集的个数）。
        return len(self.samples)

    def __getitem__(self, idx):  # 类的实例对象（p），可以像p[key]取值，当实例对象做p[key]运算时，会调用__getitem__()方法。
        return self.samples[idx]  # 返回当前具体数据，idx的最大取值为len(samples)

　数据预处理

　

　　pokemon数据集：

import torch
import os, glob
import random, csv
from torch.utils.data import Dataset, DataLoader  # Dataloader：实现batch加载数据
from torchvision import transforms
from PIL import Image
import visdom
import time


class Pokemon(Dataset):  # 继承Datastet类，如自定义模型继承Module类一样
    def __init__(self, root, resize, mode):
        super(Pokemon, self).__init__()
        self.root = root
        self.resize = resize

        self.name2lable = {}
        for name in sorted(os.listdir(os.path.join(root))):
            if not os.path.isdir(os.path.join(root, name)):
                continue
            self.name2lable[name] = len(self.name2lable.keys())

        print(self.name2lable)
        # 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%  60%-80%
            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))]
        else:  # 20% test
            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.join(self.root, filename):
            images = []
            for name in self.name2lable.keys():
                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)  # pokemon\\bulbasaur\\00000000.png
            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.name2lable[name]
                    writer.writerow([img, label, ])
                print("writer 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:  # pokemon\mewtwo\00000005.png,2
                img, label = row  # img为第一列, label为第二列
                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):  # 类的实例对象（p），可以像p[key]取值，当实例对象做p[key]运算时，会调用__getitem__()方法。
        # self.images self.labels idx:[0-len(images)]
        img, label = self.images[idx], self.labels[idx]  # 返回当前具体数据

        transf = transforms.Compose([
            lambda x: Image.open(x).convert('RGB'),  # string path ---> image data
            transforms.Resize((int(self.resize * 1.25), int(self.resize * 1.25))),
            transforms.RandomRotation(15),  # 旋转15度
            transforms.CenterCrop(self.resize),  #
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406],  # 会影响图片的可视化效果，需要做denormalize
                                 std=[0.229, 0.224, 0.225])
        ])

        img = transf(img)
        label = torch.tensor(label)
        return img, label


def main():
    viz = visdom.Visdom()
    pokemon = Pokemon('pokemon', 224, 'train')

    x, y = next(iter(pokemon))
    print('samples:', x.shape, y.shape, y)  # 打印一张图片的格式 samples: torch.Size([3, 224, 224]) torch.Size([]) tensor(2)
    # 可视化这张图片
    # viz.images(pokemon.denormalize(x), win='samples_x', opts=dict(title='sample_xx'))
    loder = DataLoader(pokemon, batch_size=32, shuffle=True)  # shuffle:每次去的batch是随机取的

    for x, y in loder:
        viz.images(pokemon.denormalize(x), nrow=8, win='batch', opts=dict(title='batch'))  # nrow:每行显示8张
        viz.text(str(y.numpy()), win='label', opts=dict(title='batch-y'))

        time.sleep(10)


if __name__ == '__main__':
    main()

　使用API完成：

from torch.utils.data import Dataset, DataLoader  
import torchvision
from torchvision import transforms
import visdom
import time


def main():
    viz = visdom.Visdom()
    transf = transforms.Compose([
        transforms.Resize((64, 64)),
        transforms.ToTensor()
    ])

    db = torchvision.datasets.ImageFolder(root='pokemon', transform=transf)
    loder = DataLoader(db, batch_size=32, shuffle=True)  # Dataloader：实现batch加载数据
    print(db.class_to_idx)  # 打印编码 {'bulbasaur': 0, 'charmander': 1, 'mewtwo': 2, 'pikachu': 3, 'squirtle': 4}

    for x, y in loder:
        viz.images(x, nrow=8, win='batch', opts=dict(title='batch'))  # nrow:每行显示8张
        viz.text(str(y.numpy()), win='label', opts=dict(title='batch-y'))
        time.sleep(10)


if __name__ == '__main__':
    main()

 2、创建模型

　Inherit from base class；Define forward graph。

import torch
from torch import nn
from torch.nn import functional as F


class ResBlk(nn.Module):
    """
    resnet block
    """

    def __init__(self, ch_in, ch_out, stride=1):
        super(ResBlk, self).__init__()

        # we add stride support for resbok, which is distinct from tutorials.
        self.conv1 = nn.Conv2d(ch_in, ch_out, kernel_size=3, stride=stride, padding=1)
        self.bn1 = nn.BatchNorm2d(ch_out)
        self.conv2 = nn.Conv2d(ch_out, ch_out, kernel_size=3, stride=1, padding=1)
        self.bn2 = nn.BatchNorm2d(ch_out)

        self.extra = nn.Sequential()
        if ch_out != ch_in:
            # [b, ch_in, h, w] => [b, ch_out, h, w]
            self.extra = nn.Sequential(
                nn.Conv2d(ch_in, ch_out, kernel_size=1, stride=stride),  # 1x1卷积核的作用
                nn.BatchNorm2d(ch_out)
            )

    def forward(self, x):
        """
        :param x: [b, ch, h, w]
        :return:
        """
        out = F.relu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))
        # short cut.
        # extra module: [b, ch_in, h, w] => [b, ch_out, h, w]
        # element-wise add:
        out = self.extra(x) + out
        out = F.relu(out)

        return out


class ResNet18(nn.Module):

    def __init__(self, num_class):
        super(ResNet18, self).__init__()

        self.conv1 = nn.Sequential(
            nn.Conv2d(3, 16, kernel_size=3, stride=3, padding=0),
            nn.BatchNorm2d(16)
        )
        # followed 4 blocks
        # [b, 64, h, w] => [b, 128, h ,w]
        self.blk1 = ResBlk(16, 32, stride=2)
        # [b, 128, h, w] => [b, 256, h, w]
        self.blk2 = ResBlk(32, 64, stride=2)
        # # [b, 256, h, w] => [b, 512, h, w]
        self.blk3 = ResBlk(64, 128, stride=2)
        # # [b, 512, h, w] => [b, 1024, h, w]
        self.blk4 = ResBlk(128, 256, stride=2)

        self.outlayer = nn.Linear(256 * 2 * 2, num_class)

    def forward(self, x):
        """

        :param x:
        :return:
        """
        x = F.relu(self.conv1(x))

        # [b, 64, h, w] => [b, 1024, h, w]
        x = self.blk1(x)
        x = self.blk2(x)
        x = self.blk3(x)
        x = self.blk4(x)

        print('after conv:', x.shape)  # [b, 512, 2, 2]
        # [b, 512, h, w] => [b, 512, 1, 1]
        # x = F.adaptive_avg_pool2d(x, [1, 1])
        # print('after pool:', x.shape)
        x = x.view(x.size(0), -1)
        x = self.outlayer(x)

        return x


def main():
    blk = ResBlk(64, 128)
    tmp = torch.randn(2, 64, 224, 224)
    out = blk(tmp)
    print('block:', out.shape)

    model = ResNet18(5)
    x = torch.randn(2, 3, 64, 64)
    out = model(x)
    print('resnet:', out.shape)

    p = sum(map(lambda p: p.numel(), model.parameters()))
    print("parameters size :", p)


if __name__ == '__main__':
    main()

3、Train & Test

　　

import torch
import os, glob
import random, csv
from torch.utils.data import Dataset, DataLoader  # Dataloader：实现batch加载数据
from torchvision import transforms
from PIL import Image
import visdom
import time


class Pokemon(Dataset):  # 继承Datastet类，如自定义模型继承Module类一样
    def __init__(self, root, resize, mode):
        super(Pokemon, self).__init__()
        self.root = root
        self.resize = resize

        self.name2lable = {}
        for name in sorted(os.listdir(os.path.join(root))):
            if not os.path.isdir(os.path.join(root, name)):
                continue
            self.name2lable[name] = len(self.name2lable.keys())

        # print(self.name2lable)
        # 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%  60%-80%
            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))]
        else:  # 20% test
            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.join(self.root, filename):
            images = []
            for name in self.name2lable.keys():
                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)  # pokemon\\bulbasaur\\00000000.png
            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.name2lable[name]
                    writer.writerow([img, label, ])
                print("writer 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:  # pokemon\mewtwo\00000005.png,2
                img, label = row  # img为第一列, label为第二列
                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):  # 类的实例对象（p），可以像p[key]取值，当实例对象做p[key]运算时，会调用__getitem__()方法。
        # self.images self.labels idx:[0-len(images)]
        img, label = self.images[idx], self.labels[idx]  # 返回当前具体数据

        transf = transforms.Compose([
            lambda x: Image.open(x).convert('RGB'),  # string path ---> image data
            transforms.Resize((int(self.resize * 1.25), int(self.resize * 1.25))),
            transforms.RandomRotation(15),  # 旋转15度
            transforms.CenterCrop(self.resize),  #
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406],  # 会影响图片的可视化效果，需要做denormalize
                                 std=[0.229, 0.224, 0.225])
        ])

        img = transf(img)
        label = torch.tensor(label)
        return img, label


def main():
    viz = visdom.Visdom()
    pokemon = Pokemon('pokemon', 224, 'train')

    x, y = next(iter(pokemon))
    print('samples:', x.shape, y.shape, y)  # 打印一张图片的格式 samples: torch.Size([3, 224, 224]) torch.Size([]) tensor(2)
    # 可视化这张图片
    # viz.images(pokemon.denormalize(x), win='samples_x', opts=dict(title='sample_xx'))
    loder = DataLoader(pokemon, batch_size=32, shuffle=True)  # shuffle:每次去的batch是随机取的

    for x, y in loder:
        viz.images(pokemon.denormalize(x), nrow=8, win='batch', opts=dict(title='batch'))  # nrow:每行显示8张
        viz.text(str(y.numpy()), win='label', opts=dict(title='batch-y'))

        time.sleep(10)


if __name__ == '__main__':
    main()

import torch
from torch import nn
from torch.nn import functional as F


class ResBlk(nn.Module):
    """
    resnet block
    """

    def __init__(self, ch_in, ch_out, stride=1):
        super(ResBlk, self).__init__()

        # we add stride support for resbok, which is distinct from tutorials.
        self.conv1 = nn.Conv2d(ch_in, ch_out, kernel_size=3, stride=stride, padding=1)
        self.bn1 = nn.BatchNorm2d(ch_out)
        self.conv2 = nn.Conv2d(ch_out, ch_out, kernel_size=3, stride=1, padding=1)
        self.bn2 = nn.BatchNorm2d(ch_out)

        self.extra = nn.Sequential()
        if ch_out != ch_in:
            # [b, ch_in, h, w] => [b, ch_out, h, w]
            self.extra = nn.Sequential(
                nn.Conv2d(ch_in, ch_out, kernel_size=1, stride=stride),  # 1x1卷积核的作用
                nn.BatchNorm2d(ch_out)
            )

    def forward(self, x):
        """
        :param x: [b, ch, h, w]
        :return:
        """
        out = F.relu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))
        # short cut.
        # extra module: [b, ch_in, h, w] => [b, ch_out, h, w]
        # element-wise add:
        out = self.extra(x) + out
        out = F.relu(out)

        return out


class ResNet18(nn.Module):

    def __init__(self, num_class):
        super(ResNet18, self).__init__()

        self.conv1 = nn.Sequential(
            nn.Conv2d(3, 16, kernel_size=3, stride=3, padding=0),
            nn.BatchNorm2d(16)
        )
        # followed 4 blocks
        # [b, 64, h, w] => [b, 128, h ,w]
        self.blk1 = ResBlk(16, 32, stride=3)
        # [b, 128, h, w] => [b, 256, h, w]
        self.blk2 = ResBlk(32, 64, stride=3)
        # # [b, 256, h, w] => [b, 512, h, w]
        self.blk3 = ResBlk(64, 128, stride=2)
        # # [b, 512, h, w] => [b, 1024, h, w]
        self.blk4 = ResBlk(128, 256, stride=2)

        self.outlayer = nn.Linear(256 * 3 * 3, num_class)

    def forward(self, x):
        """

        :param x:
        :return:
        """
        x = F.relu(self.conv1(x))

        # [b, 64, h, w] => [b, 1024, h, w]
        x = self.blk1(x)
        x = self.blk2(x)
        x = self.blk3(x)
        x = self.blk4(x)

        # print('after conv:', x.shape)  # [b, 512, 2, 2]
        # [b, 512, h, w] => [b, 512, 1, 1]
        # x = F.adaptive_avg_pool2d(x, [1, 1])
        # print('after pool:', x.shape)
        x = x.view(x.size(0), -1)
        x = self.outlayer(x)

        return x


def main():
    blk = ResBlk(64, 128)
    tmp = torch.randn(2, 64, 224, 224)
    out = blk(tmp)
    print('block:', out.shape)

    model = ResNet18(5)
    x = torch.randn(2, 3, 64, 64)
    out = model(x)
    print('resnet:', out.shape)  # resnet: torch.Size([2, 5]) batch:2

    p = sum(map(lambda p: p.numel(), model.parameters()))
    print("parameters size :", p)


if __name__ == '__main__':
    main()

import torch
from torch import optim, nn
import visdom
import torchvision
from torch.utils.data import Dataset, DataLoader
from wupiao import Pokemon
from Mymodel import ResNet18

batchsz = 32
lr = 1e-3
epochs = 10
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_loder = DataLoader(train_db, batch_size=batchsz, shuffle=True, num_workers=4)
val_loder = DataLoader(val_db, batch_size=batchsz, num_workers=2)
test_loder = DataLoader(test_db, batch_size=batchsz, num_workers=2)

viz = visdom.Visdom()

def evalute(model, loder):  # 对validation() and test()是相同的操作
    correct = 0
    total = len(loder.dataset)
    for x,y in loder:
        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)
    optimizer = optim.Adam(model.parameters(), lr=lr)
    criterion = nn.CrossEntropyLoss()

    best_acc, best_epoch = 0, 0
    global_step = 0
    viz.line([0], [-1], win='loss', opts=dict(title='loss'))
    viz.line([0], [-1], win='val_acc', opts=dict(title='val_acc'))
    for epoch in range(epochs):
        for step, (x, y) in enumerate(train_loder):
            # x:[b, 3, 224, 224] y:[b]

            logits = model(x)
            loss = criterion(logits, y)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            viz.line([loss.item()], [global_step], win='loss', update='append')
            global_step += 1
        if epochs % 2 == 0:  # 做一个卷积度测试
            val_acc = evalute(model, val_loder)
            if val_acc > best_acc:
                best_acc = val_acc
                best_epoch = epoch

                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)
    model.load_state_dict(torch.load('best.mdl'))  # 用最好的模型覆盖之前的模型
    print('loaded from ckpt!')

    test_acc = evalute(model, test_loder)  # 使用最好的model来测试
    print('test acc:', test_acc)



if __name__ == '__main__':
    main()

4、迁移学习（Transfer learning）

 　　

　　　　

import torch
from torch import nn
from matplotlib import pyplot as plt


class Flatten(nn.Module):
    def __init__(self):
        super(Flatten, self).__init__()

    def forward(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):  # 6中照片类型
        plt.subplot(2, 2, i + 1)
        plt.tight_layout()
        plt.imshow(img[i][0] * 0.308 + 0.1307, cmap='gray', interpolation='none')
        plt.title("{} : {}".format(name, label[i].item()))
        plt.xticks()
        plt.yticks()
    plt.show()

import torch
from torch import optim, nn
import visdom
import torchvision
from torch.utils.data import Dataset, DataLoader
from wupiao import Pokemon
# from Mymodel import ResNet18
from  torchvision.models import resnet18
from utils import Flatten
batchsz = 32
lr = 1e-3
epochs = 10
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_loder = DataLoader(train_db, batch_size=batchsz, shuffle=True, num_workers=4)
val_loder = DataLoader(val_db, batch_size=batchsz, num_workers=2)
test_loder = DataLoader(test_db, batch_size=batchsz, num_workers=2)

viz = visdom.Visdom()

def evalute(model, loder):  # 对validation() and test()是相同的操作
    correct = 0
    total = len(loder.dataset)
    for x,y in loder:
        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)
    train_model = resnet18(pretrained=True)
    model = nn.Sequential(*list(train_model.children())[:-1],
                          Flatten(), # ([b, 512, 1, 1]) => ([b, 512])
                          nn.Linear(512, 5)
                          )
    # x = torch.randn(2, 3, 224, 224)
    # print(model(x).shape)  # torch.Size([2, 512, 1, 1]) =>torch.Size([2, 5])

    optimizer = optim.Adam(model.parameters(), lr=lr)
    criterion = nn.CrossEntropyLoss()

    best_acc, best_epoch = 0, 0
    global_step = 0
    viz.line([0], [-1], win='loss', opts=dict(title='loss'))
    viz.line([0], [-1], win='val_acc', opts=dict(title='val_acc'))
    for epoch in range(epochs):
        for step, (x, y) in enumerate(train_loder):
            # x:[b, 3, 224, 224] y:[b]

            logits = model(x)
            loss = criterion(logits, y)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            viz.line([loss.item()], [global_step], win='loss', update='append')
            global_step += 1
        if epochs % 2 == 0:  # 做一个卷积度测试
            val_acc = evalute(model, val_loder)
            if val_acc > best_acc:
                best_acc = val_acc
                best_epoch = epoch

                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)
    model.load_state_dict(torch.load('best.mdl'))  # 用最好的模型覆盖之前的模型
    print('loaded from ckpt!')

    test_acc = evalute(model, test_loder)  # 使用最好的model来测试
    print('test acc:', test_acc)



if __name__ == '__main__':
    main()

---