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()