pytorch的torchvision内置了resnet主干网络,想要训练自定义分类数据集,只需要将全连接层的替换即可实现自定义数据集分类

本次我们使用的是resnet18 做自定义数据集分类

项目依赖

numpy1.20.3
opencv-contrib-python4.5.3.56
opencv-python4.5.1.48
opencv-python-headless4.5.1.48
Pillow8.2.0
tensorboard2.4.1
tensorboard-plugin-wit1.8.0
torch1.7.1
torchvision0.8.2
tqdm4.60.0

数据集组织形式

我们将需要分类的图像数据放到某个目录下, 每个文件夹代表一个类别(类别使用数值代表)

如下图所示:

数据集加载

import glob
import os
from PIL import Image


def default_loader(path):
    with open(path, 'rb') as f:
        with Image.open(f) as img:
            return img.convert('RGB')


class Dataset:

    def __init__(self, image_root_path, data_transforms=None, image_format='png'):
        self.data_transforms = data_transforms
        self.image_root_path = image_root_path
        self.image_format = image_format
        self.images = []
        self.labels = []
        classes_folders = os.listdir(self.image_root_path)
        for cls_folder in classes_folders:
            folder_path = os.path.join(self.image_root_path, cls_folder)
            if os.path.isdir(folder_path):
                images_path = os.path.join(folder_path, "*.{}".format(self.image_format))
                images = glob.glob(images_path)
                self.images.extend(images)

    def __len__(self):
        return len(self.images)

    def __getitem__(self, item):
        image_file = self.images[item]
        label_name = os.path.basename(os.path.dirname(image_file))
        image = default_loader(image_file)
        if self.data_transforms is not None:
            image = self.data_transforms(image)

        return image, int(label_name)

模型训练

import config
import torch
import torch.optim as optim
from torch.optim import lr_scheduler
from torch import nn
import os
from tqdm import tqdm
from torchvision import models, transforms
from torch.utils.tensorboard import SummaryWriter
from dataset import Dataset


def train(model, loss_func, dataset, optimizer, epoch, writer):
    model.train()
    batch_loss = 0
    item = 0
    for batch, (image, label) in tqdm(enumerate(dataset)):
        image = image.to(config.device)
        label = label.to(config.device)
        optimizer.zero_grad()
        output = model(image)
        _, pred = torch.max(output, 1)
        loss = loss_func(output, label)
        loss.backward()
        optimizer.step()
        writer.add_images("train_images", image, epoch)
        writer.add_scalar("train_loss", loss, epoch)
        print("Train Epoch = {} Loss = {}".format(epoch, loss.data.item()))
        batch_loss += loss.data.item()
        item += 1

    return batch_loss / item


def valid(model, loss_func, dataset, epoch, writer):
    model.eval()
    batch_loss = 0
    item = 0
    with torch.no_grad():
        for batch, (image, label) in tqdm(enumerate(dataset)):
            image = image.to(config.device)
            label = label.to(config.device)
            output = model(image)
            loss = loss_func(output, label)
            writer.add_images("valid_images", image, epoch)
            writer.add_scalar("valid_loss", loss, epoch)
            batch_loss += loss.data.item()
            item += 1
            print("Valid Epoch = {} Loss = {}".format(epoch, loss.data.item()))
    return batch_loss / item


def train_model(model, loss_func, optimizer, step_scheduler, num_epochs=config.epoch):
    train_transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])  # 各通道颜色的均值和方差,用于归一化
    ])
    valid_transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])  # 各通道颜色的均值和方差,用于归一化
    ])
    train_dataset = Dataset(config.train_image_path, train_transform, config.image_format)
    valid_dataset = Dataset(config.valid_image_path, valid_transform, config.image_format)

    train_dataloader = torch.utils.data.DataLoader(
        train_dataset, batch_size=config.batch_size,
        shuffle=True, num_workers=config.num_workers,
    )
    valid_dataloader = torch.utils.data.DataLoader(
        valid_dataset, batch_size=config.batch_size,
        shuffle=True, num_workers=config.num_workers
    )
    start_epoch = 0
    # 断点继续训练
    if config.resume:
        checkpoint = torch.load(config.chkpt)  # 加载断点
        model.load_state_dict(checkpoint['net'])  # 加载模型可学习参数
        optimizer.load_state_dict(checkpoint['optimizer'])  # 加载优化器参数
        start_epoch = checkpoint['epoch']  # 设置开始的epoch
    writer = SummaryWriter(log_dir="./runs")
    # images, _ = next(iter(train_dataloader))
    # writer.add_graph(model, images)
    for epoch in range(start_epoch + 1, num_epochs):
        train_epoch_loss = train(model, loss_func, train_dataloader, optimizer, epoch, writer)
        valid_epoch_loss = valid(model, loss_func, valid_dataloader, epoch, writer)
        step_scheduler.step()
        # 模型保存
        if epoch % config.save_model_iter == 0:
            checkpoint = {
                "net": model.state_dict(),
                'optimizer': optimizer.state_dict(),
                "epoch": epoch
            }
            save_model_file = os.path.join(config.model_output_dir, "epoch_{}.pth".format(epoch))
            if not os.path.exists(config.model_output_dir):
                os.makedirs(config.model_output_dir)
            torch.save(checkpoint, save_model_file)
        if train_epoch_loss < config.best_loss or valid_epoch_loss < config.best_loss:
            checkpoint = {
                "net": model.state_dict(),
                'optimizer': optimizer.state_dict(),
                "epoch": epoch
            }
            save_model_file = os.path.join(config.model_output_dir, "best_{}.pth".format(epoch))
            if not os.path.exists(config.model_output_dir):
                os.makedirs(config.model_output_dir)
            torch.save(checkpoint, save_model_file)
        if epoch % 10 == 0:
            print("Epoch = {} Train Loss = {} Valid Loss = {}".format(epoch, train_epoch_loss, valid_epoch_loss))
    writer.close()


if __name__ == '__main__':
    backbone = models.resnet18(pretrained=True)
    num_fits = backbone.fc.in_features
    backbone.fc = nn.Linear(num_fits, config.num_classes)  # 替换最后一个全连接层
    model_ft = backbone.to(config.device)
    criterion = nn.CrossEntropyLoss()
    optimizer_ft = optim.Adam(model_ft.parameters(), lr=config.lr)
    scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=10, gamma=0.1)
    train_model(model_ft, criterion, optimizer_ft, scheduler, config.epoch)

模型预测

import glob
import os
import cv2
import config
import torch
import numpy as np
from torch import nn
from PIL import Image
from torchvision import models
import torchvision.transforms as transforms

transform_test = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
])


def predict_images(image_file, label, model):
    image = Image.open(image_file)
    image = image.convert("RGB")
    numpy_array = np.asarray(image.copy())
    image = transform_test(image)
    image = image.unsqueeze_(0).to(config.device)
    with torch.no_grad():
        outputs = model(image)
        outputs = outputs.to('cpu')
    predict_label = torch.max(outputs, dim=1)[1].data.numpy()[0]
    if predict_label != label:
        print("predict error image = {}".format(image_file))

    print("测试类别={}".format(predict_label))
    cv2.imshow("image", numpy_array)
    cv2.waitKey(0)


def get_image_label_to_predict():
    model = models.resnet18(pretrained=False)
    num_fits = model.fc.in_features
    model.fc = nn.Linear(num_fits, config.num_classes)
    model.load_state_dict(torch.load(config.predict_model)['net'])
    model.eval()
    model.to(config.device)
    classes_dir = os.listdir(config.predict_image_path)
    for label in classes_dir:
        label_path = os.path.join(config.predict_image_path, label)
        if os.path.isdir(label_path):
            images = glob.glob(os.path.join(label_path, "*.{}".format(config.image_format)))
            for img in images:
                predict_images(img, int(label), model)


if __name__ == '__main__':
    get_image_label_to_predict()

完整代码项目地址

Github 地址 https://github.com/pythondever/pytorch_resnet18_image_classify

如果您觉得这个项目对您有帮助,欢迎 star