使用resnet50简单搭建我们自己的分类模型

dataset

from PIL import Image
from torch.utils.data import Dataset
import os

class MyDataLoader(Dataset):
    def __init__(self, root_dir, transform, augment=True):
        self.root_dir = root_dir
        self.augment = augment
        self.testNameToNum = {"ok": 0, "wrong": 1}
        self.numToName = {0: "良品", 1: "连锡"}
        self.trainNameToNum = {"良品": 0, "连锡": 1}
        # 判断文件夹dirname是否存在
        if not os.path.exists(self.root_dir):
            print("error: folder \"", self.root_dir, "\" not exits!")
            raise Exception("error: folder" + self.root_dir + "文件夹" + "not exits!")

        self.names = os.listdir(self.root_dir)
        # 创建5*2的数据集
        self.images_path = self.get_all_images_path()
        # 5个数据的标签
        self.images_labels = self.get_all_images_labels()

        self.images_path = sum(self.images_path, [])
        self.transform = transform

    def get_all_images_path(self):
        all_images_path = []
        for root, dirs, files in os.walk(self.root_dir):
            every_dir_images_path = []
            for file in files:
                path = os.path.join(root, file)
                every_dir_images_path.append(path)
            all_images_path.append(every_dir_images_path)
        return all_images_path[1:]

    def get_all_images_labels(self):
        all_images_labels = []
        for i in range(len(self.names)):
            every_dir_images_path = self.images_path[i]
            every_dir_images_labels = []
            if self.names[i] in self.testNameToNum:
                every_dir_images_labels = [self.testNameToNum[self.names[i]]] * len(every_dir_images_path)
            elif self.names[i] in self.trainNameToNum:
                every_dir_images_labels = [self.trainNameToNum[self.names[i]]] * len(every_dir_images_path)
            if every_dir_images_labels is None:
                print("error: folder \"", self.root_dir, "\" not exits!")
                raise Exception(self.names[i] + "没有符合的标签")
            all_images_labels.append(every_dir_images_labels)
        return sum(all_images_labels, [])

    # 根据索引获取data和label
    def __getitem__(self, index):
        ig_path = self.images_path[index]
        label = self.images_labels[index]
        ######################################
        pil_image = Image.open(ig_path).convert('RGB')
        #########################################
        data = self.transform(pil_image)
        return data, label

    # 获取数据集的大小
    def __len__(self):
        return len(self.images_path)

# if __name__ == '__main__':
#     dirname = "data//test"
#     data = MyDataLoader(dirname)
#     print(f'data size is : {len(data)}')
#
#     image,label = data[1]
#     print(str(data))

train

import torch
from torchvision import datasets, models, transforms
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
import time

import numpy as np
import matplotlib.pyplot as plt
import os
from tqdm import tqdm
from MyDataLoader import MyDataLoader

transform = transforms.Compose([
    transforms.RandomResizedCrop(size=256, scale=(0.8, 1.0)),
    transforms.RandomRotation(degrees=15),
    transforms.RandomHorizontalFlip(),
    transforms.CenterCrop(size=224),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406],
                         [0.229, 0.224, 0.225])
])

# 三、加载数据
# torchvision.transforms包DataLoader是 Pytorch 重要的特性，它们使得数据增加和加载数据变得非常简单。
# 使用 DataLoader 加载数据的时候就会将之前定义的数据 transform 就会应用的数据上了。
dataset = 'data'
train_directory = os.path.join(dataset, 'train')
test_directory = os.path.join(dataset, 'test')

train_dataset = MyDataLoader(train_directory, transform)
test_dataset = MyDataLoader(test_directory, transform)

train_data_size = len(train_dataset)
test_data_size = len(test_dataset)

batch_size = 64
num_classes = 6

train_data = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_data = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

# 四、迁移学习
# 这里使用ResNet-50的预训练模型。
resnet50 = models.resnet50(pretrained=True)

# 在PyTorch中加载模型时，所有参数的‘requires_grad’字段默认设置为true。
# 这意味着对参数值的每一次更改都将被存储，以便在用于训练的反向传播图中使用。
# 这增加了内存需求。由于预训练的模型中的大多数参数已经训练好了，因此将requires_grad字段重置为false。
for param in resnet50.parameters():
    param.requires_grad = False

# 为了适应自己的数据集，将ResNet-50的最后一层替换为，
# 将原来最后一个全连接层的输入喂给一个有256个输出单元的线性层，接着再连接ReLU层和Dropout层，
# 然后是256 x 6的线性层，输出为6通道的softmax层。
fc_inputs = resnet50.fc.in_features
resnet50.fc = nn.Sequential(
    nn.Linear(fc_inputs, 256),
    nn.ReLU(),
    nn.Dropout(0.4),
    nn.Linear(256, 2),
    nn.LogSoftmax(dim=1)
)

# 用GPU进行训练。
resnet50 = resnet50.to('cpu')

# 定义损失函数和优化器。
loss_func = nn.CrossEntropyLoss()
optimizer = optim.Adam(resnet50.parameters())


# 五、训练
def train_and_valid(model, loss_function, optimizer, epochs=25):
    device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu")
    history = []
    best_acc = 0.0
    best_epoch = 0

    for epoch in range(epochs):
        epoch_start = time.time()
        print("Epoch: {}/{}".format(epoch + 1, epochs))

        model.train()

        train_loss = 0.0
        train_acc = 0.0
        valid_loss = 0.0
        valid_acc = 0.0

        for i, (inputs, labels) in enumerate(tqdm(train_data)):
            inputs = inputs.to(device)
            labels = labels.to(device)

            # 因为这里梯度是累加的，所以每次记得清零
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = loss_function(outputs, labels)
            loss.backward()
            optimizer.step()
            train_loss += loss.item() * inputs.size(0)
            ret, predictions = torch.max(outputs.data, 1)
            correct_counts = predictions.eq(labels.data.view_as(predictions))
            acc = torch.mean(correct_counts.type(torch.FloatTensor))
            train_acc += acc.item() * inputs.size(0)

        with torch.no_grad():
            model.eval()

            for j, (inputs, labels) in enumerate(tqdm(test_data)):
                inputs = inputs.to(device)
                labels = labels.to(device)
                outputs = model(inputs)
                loss = loss_function(outputs, labels)
                valid_loss += loss.item() * inputs.size(0)
                ret, predictions = torch.max(outputs.data, 1)
                correct_counts = predictions.eq(labels.data.view_as(predictions))
                acc = torch.mean(correct_counts.type(torch.FloatTensor))
                valid_acc += acc.item() * inputs.size(0)

        avg_train_loss = train_loss / train_data_size
        avg_train_acc = train_acc / train_data_size

        avg_valid_loss = valid_loss / test_data_size
        avg_valid_acc = valid_acc / test_data_size

        history.append([avg_train_loss, avg_valid_loss, avg_train_acc, avg_valid_acc])

        if best_acc < avg_valid_acc:
            best_acc = avg_valid_acc
            best_epoch = epoch + 1

        epoch_end = time.time()

        print(
            "Epoch: {:03d}, Training: Loss: {:.4f}, Accuracy: {:.4f}%, \n\t\tValidation: Loss: {:.4f}, Accuracy: {:.4f}%, Time: {:.4f}s".format(
                epoch + 1, avg_valid_loss, avg_train_acc * 100, avg_valid_loss, avg_valid_acc * 100,
                epoch_end - epoch_start
            ))
        print("Best Accuracy for validation : {:.4f} at epoch {:03d}".format(best_acc, best_epoch))
        MODEL_SAVE_PATH = "./"
        MODEL_NAME = 'models/'+dataset+'_model_'+str(epoch+1)+'.pt'
        torch.save(model, os.path.join(MODEL_SAVE_PATH, MODEL_NAME))
    return model, history


num_epochs = 30
trained_model, history = train_and_valid(resnet50, loss_func, optimizer, num_epochs)
torch.save(history, 'models/' + dataset + '_history.pt')

history = np.array(history)
plt.plot(history[:, 0:2])
plt.legend(['Tr Loss', 'Val Loss'])
plt.xlabel('Epoch Number')
plt.ylabel('Loss')
plt.ylim(0, 1)
plt.savefig(dataset + '_loss_curve.png')
plt.show()

plt.plot(history[:, 2:4])
plt.legend(['Tr Accuracy', 'Val Accuracy'])
plt.xlabel('Epoch Number')
plt.ylabel('Accuracy')
plt.ylim(0, 1)
plt.savefig(dataset + '_accuracy_curve.png')
plt.show()