Pytorch-Resnet迁移学习记录

本文介绍了如何使用预训练的ResNet50模型进行迁移学习,针对'cut_image_classify'和'cut_image_test'数据集,通过调整模型结构和优化器,进行训练和验证,最终保存最佳模型并绘制损失和准确率曲线。

摘要生成于 C知道 ,由 DeepSeek-R1 满血版支持, 前往体验 >

暂时不解释,仅作为记录:

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

image_transforms = {
    'cut_image_classify': 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])
    ]),
    'cut_image_test': transforms.Compose([
        transforms.Resize(size=256),
        transforms.CenterCrop(size=224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406],
                             [0.229, 0.224, 0.225])
    ])
}

dataset = 'datasets'
# train_directory = os.path.join(dataset, 'train')
train_directory = os.path.join(dataset, 'cut_image_classify')
# valid_directory = os.path.join(dataset, 'valid')
valid_directory = os.path.join(dataset, 'cut_image_test')

batch_size = 4
num_classes = 5

data = {
    'cut_image_classify': datasets.ImageFolder(root=train_directory, transform=image_transforms['cut_image_classify']),
    'cut_image_test': datasets.ImageFolder(root=valid_directory, transform=image_transforms['cut_image_test'])
}

train_data_size = len(data['cut_image_classify'])
valid_data_size = len(data['cut_image_test'])

train_data = DataLoader(data['cut_image_classify'], batch_size=batch_size, shuffle=True)
valid_data = DataLoader(data['cut_image_test'], batch_size=batch_size, shuffle=True)

print(train_data_size, valid_data_size)

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

for param in ResNet50.parameters():
    param.requires_grad = False

fc_inputs = ResNet50.fc.in_features
ResNet50.fc = nn.Sequential(
    nn.Linear(fc_inputs, 256),
    nn.ReLU(),
    nn.Dropout(0.4),
    nn.Linear(256, 5),
    nn.LogSoftmax(dim=1)
)
# # 用GPU进行训练
# resnet50 = resnet50.to('cuda:0')

# 用CPU训练
resnet50 = ResNet50.to('cpu')
# loss_func = nn.CrossEntropyLoss()
loss_func = nn.NLLLoss()
optimizer = optim.Adam(resnet50.parameters(), lr=0.0002)


def train_and_valid(model, loss_function, optimizer, epochs=25):
    device = torch.device("cuda:0" 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(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(valid_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 / valid_data_size
        avg_valid_acc = valid_acc / valid_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
            torch.save(model, 'models/' + dataset + '_model_' + str(epoch + 1) + '.pt')
        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_train_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))

        # torch.save(model, 'models/' + dataset + '_model_' + str(epoch + 1) + '.pt')
    return model, history

if __name__ == "__main__":
    num_epochs = 600
    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, 2)
    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()
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值