resnet18训练自定义数据集

目录结构

dogsData.py

import json

import torch
import os, glob
import random, csv

from PIL import Image
from torch.utils.data import Dataset, DataLoader

from torchvision import transforms
from torchvision.transforms import InterpolationMode

class Dogs(Dataset):

    def __init__(self, root, resize, mode):
        super().__init__()
        self.root = root
        self.resize = resize
        self.nameLable = {}
        for name in sorted(os.listdir(os.path.join(root))):
            if not os.path.isdir(os.path.join(root, name)):
                continue
            self.nameLable[name] = len(self.nameLable.keys())

        if not os.path.exists(os.path.join(self.root, 'label.txt')):
            with open(os.path.join(self.root, 'label.txt'), 'w', encoding='utf-8') as f:
                f.write(json.dumps(self.nameLable, ensure_ascii=False))

        # print(self.nameLable)
        self.images, self.labels = self.load_csv('images.csv')
        # print(self.labels)

        if mode == 'train':
            self.images = self.images[:int(0.8*len(self.images))]
            self.labels = self.labels[:int(0.8*len(self.labels))]
        elif mode == 'val':
            self.images = self.images[int(0.8*len(self.images)):int(0.9*len(self.images))]
            self.labels = self.labels[int(0.8*len(self.labels)):int(0.9*len(self.labels))]
        else:
            self.images = self.images[int(0.9*len(self.images)):]
            self.labels = self.labels[int(0.9*len(self.labels)):]

    def load_csv(self, filename):

        if not os.path.exists(os.path.join(self.root, filename)):
            images = []
            for name in self.nameLable.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))

            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.nameLable[name]
                    writer.writerow([img, label])
            print('csv write succesful')

        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 denormalize(self, x_hat):
        mean = [0.485, 0.456, 0.406]
        std = [0.229, 0.224, 0.225]
        # x_hot = (x-mean)/std
        # x = x_hat * std = mean
        # x : [c, w, h]
        # 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 __len__(self):
        return len(self.images)

    def __getitem__(self, idx):
        # print(idx, len(self.images), len(self.labels))
        img, label = self.images[idx], self.labels[idx]

        # 将字符串路径转换为tensor数据
        # print(self.resize, type(self.resize))
        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),
            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():

    import visdom
    import time

    viz = visdom.Visdom()

    # func1 通用
    db = Dogs('Images_Data_Dog', 224, 'train')
    # 取一张
    # x,y = next(iter(db))
    # print(x.shape, y)
    # viz.image(db.denormalize(x), win='sample_x', opts=dict(title='sample_x'))

    # 取一个batch
    loader = DataLoader(db, batch_size=32, shuffle=True, num_workers=8)
    print(len(loader))
    print(db.nameLable)
    # for x, y in loader:
    #     # print(x.shape, y)
    #     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)

    # # fun2
    # import torchvision
    # tf = transforms.Compose([
    #     transforms.Resize((64, 64)),
    #     transforms.RandomRotation(15),
    #     transforms.ToTensor(),
    # ])
    # db = torchvision.datasets.ImageFolder(root='pokemon', transform=tf)
    # loader = DataLoader(db, batch_size=32, shuffle=True)
    # print(len(loader))
    # for x, y in loader:
    #     # print(x.shape, y)
    #     viz.images(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

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

    def forward(self, x):
        shape = torch.prod(torch.tensor(x.shape[1:])).item()
        return x.view(-1, shape)

train.py

import os
import sys
base_path = os.path.dirname(os.path.abspath(__file__))
sys.path.append(base_path)
base_path = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.append(base_path)
import torch
import visdom
from torch import optim, nn
import torchvision

from torch.utils.data import DataLoader

from dogs_train.utils import Flatten
from dogsData import Dogs

from torchvision.models import resnet18

viz = visdom.Visdom()

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

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

train_db = Dogs('Images_Data_Dog', 224, mode='train')
val_db = Dogs('Images_Data_Dog', 224, mode='val')
test_db = Dogs('Images_Data_Dog', 224, mode='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)

def evalute(model, loader):
    correct = 0
    total = len(loader.dataset)
    for x, y in loader:
        x = x.to(device)
        y = y.to(device)
        with torch.no_grad():
            logist = model(x)
            pred = logist.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],
                          Flatten(),  # [b, 512, 1, 1] => [b, 512]
                          nn.Linear(512, 27)
                          ).to(device)

    x = torch.randn(2, 3, 224, 224).to(device)
    print(model(x).shape)

    optimizer = optim.Adam(model.parameters(), lr=lr)
    criteon = 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 = x.to(device)
            y = 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
        if epoch % 2 == 0:
            val_acc = evalute(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)

    model.load_state_dict(torch.load('best.mdl'))
    print('loader from ckpt')

    test_acc = evalute(model, test_loader)
    print(test_acc)

if __name__ == '__main__':
    main()