关于pokemon的深度学习

首先声明：这是一个半成品，训练到train_scratch时发现结果输出错误，具体原因困扰几天，暂时先放入csdn，以后有时间 \ 涨姿势了再回看这个bug。

这个训练需要一个数据包visdom，下次训练不要忘记这个东西呦！

---

首先是对数据集图像的处理代码

这块代码是没问题的！！！！！！！！！！！！！！！

import torch
import os, glob
import random, csv

from torch.utils.data import Dataset, DataLoader

from torchvision import transforms
from PIL import Image



class Pokemon(Dataset):
    def __init__(self, root, resize, model):
        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)

        #image, label

        self.images, self.labels = self.load_csv('image.csv')

        if model == 'train':  #取百分之六十作为train
            self.images = self.images[:int(0.6*len(self.images))]
            self.labels = self.labels[:int(0.6*len(self.labels))]
        elif model == 'val':  #取百分之二十作为,从百分之六十到百分之八十
            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.6*len(self.labels))]
        else:                #取百分之二十作为text
            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\\newtwo\\00001.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'))
                images += glob.glob(os.path.join(self.root, name, '*.gif'))

            #1168张图片，'pokemonimg\\bulbasaur\\00000000.png'
            print(len(images), images)

            random.shuffle(images)
            with open(os.path.join(self.root, filename), mode='w', newline='')as f:
                writer = csv.writer(f)
                for img in images:  #'pokemon\\bulbasaur\\00000000.png'
                    name = img.split(os.sep)[-2]
                    label = self.name2label[name]
                    writer.writerow([img, label])#存储图片路径和标签
                print('writen into csv file:', filename)

        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)

############为了可视化，将标准化的x还原
    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_hot*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):
        img, label = self.images[idx], self.labels[idx]

        tf = transforms.Compose([
            lambda x:Image.open(x).convert('RGB'),
            transforms.Resize((int(self.resize * 1.25), int(self.resize * 1.25))),  #调整图像的大小
            transforms.RandomRotation(15),   #随机旋转，15是指正负旋转15度
            transforms.CenterCrop(self.resize),  #中心裁剪
            transforms.ToTensor(),   #将PIL的image转化文张量
            transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                 std=[0.229, 0.224, 0.225])  #这些数据来自imagenet数据集，不用管，就用可以了
        ])

        img = tf(img)
        label = torch.tensor(label)

        return img, label

def main():

    import visdom   #可视化
    import time

    viz = visdom.Visdom()

    db = Pokemon('pokemon', 64, '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'))

    loader = DataLoader(db, batch_size=32, shuffle=True, num_workers=8)  ##如果一次性加载图片过多就使用num—works=8使得多线程操作
    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()

 然后是resnet网络代码

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):
        """
        :param ch_in:
        :param ch_out:
        """
        super(ResBlk, self).__init__()

        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),
                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, 16, h, w] => [b, 32, h ,w]
        self.blk1 = ResBlk(16, 32, stride=3)
        # [b, 32, h, w] => [b, 64, h, w]
        self.blk2 = ResBlk(32, 64, stride=3)
        # # [b, 64, h, w] => [b, 128, h, w]
        self.blk3 = ResBlk(64, 128, stride=2)
        # # [b, 128, h, w] => [b, 256, h, w]
        self.blk4 = ResBlk(128, 256, stride=2)

        # [b, 256, 7, 7]
        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(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)
    tmp = torch.randn(2, 3, 224, 224)
    out = model(tmp)
    print('resnet:', out.shape)

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


if __name__ == '__main__':
    main()

最后就是训练的代码

本来这个段代码运行后会出现loss值和loss图，但是不知道为什么一直没有调出来！！！！！！！！！！！！！！！！！！！！！！！！！！难受！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！

import torch
from torch import optim, nn
import visdom
from torch.utils.data import DataLoader, Dataset

from pokemon import Pokemon
from resnet import ResNet18


batchsz = 32
lr = 1e-3
epochs = 10

device = torch.device('cuda')
torch.manual_seed(1234)

train_db = Pokemon('pokemon', 224, model='train')
val_db = Pokemon('pokemon', 224, model='val')
test_db = Pokemon('pokemon', 224, model='test')
train_loader = DataLoader(train_db, batch_size=batchsz, shuffle=True, num_workers=4)
val_loader = DataLoader(val_db, batch_size=batchsz, num_workers=2)
test_loader = DataLoader(test_db, batch_size=batchsz, num_workers=2)

viz = visdom.Visdom()

def evaluate(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():
            output = model(x)
            pred = output.argmax(dim=1)
            correct += torch.eq(pred, y).sum().item()
    return correct/total



def main():
    model = ResNet18(5).to(device)
    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_loader):
            # x: [b, 3, 224, 224], y: [b]
            x, y = x.to(device), y.to(device)
            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 epoch % 2 == 0:
                val_acc = evaluate(model, val_loader)
                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+1)
        model.load_state_dict(torch.load('best.mdl'))
        print('loaded from ckpt!')
        test_acc = evaluate(model, test_loader)
        print('test acc:', test_acc)

OK，就这么多，数据集可以在B站找到连接，也可以从github获取。