乐趣区

关于python:使用PyTorch实现迁移学习

在本文中,咱们将介绍如何应用 Python 的 PyTorch 库进行迁徙学习。迁徙学习是一种机器学习技术,它容许咱们利用一个事后训练好的模型对新工作进行训练,从而节俭了大量的计算资源和工夫。

1. 装置并导入库

首先,咱们须要装置 PyTorch 库。在终端或命令提示符中输出以下命令进行装置:

pip install torch torchvision

而后导入所需的库:

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
from torchvision import datasets, models, transforms

2. 数据预处理

在本例中,咱们将应用 CIFAR-10 数据集进行迁徙学习。首先,咱们须要对数据进行预处理:

transform = transforms.Compose([transforms.RandomHorizontalFlip(),
    transforms.RandomCrop(32, padding=4),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

trainset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=100, shuffle=True, num_workers=2)

testset = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=100, shuffle=False, num_workers=2)

classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

3. 加载预训练模型

接下来,咱们将加载一个预训练的 ResNet-18 模型。预训练模型在 ImageNet 数据集上进行了训练,因而它曾经具备很好的特征提取能力。

model = models.resnet18(pretrained=True)

4. 批改模型的全连贯层

因为咱们的工作是对 CIFAR-10 数据集进行分类,因而咱们须要将模型的全连贯层(输入层)批改为具备 10 个输入单元的新层。咱们能够通过以下代码实现这一指标:

num_features = model.fc.in_features
model.fc = nn.Linear(num_features, 10)

5. 训练模型

咱们将应用穿插熵损失和随机梯度降落(SGD)优化器进行训练。接下来,咱们将定义训练函数:

def train_model(model, dataloader, criterion, optimizer, num_epochs=25):
    for epoch in range(num_epochs):
        running_loss = 0.0
        running_corrects = 0
        
        for inputs, labels in dataloader:
            optimizer.zero_grad()

            outputs = model(inputs)
            _, preds = torch.max(outputs, 1)
            loss = criterion(outputs, labels)

            loss.backward()
            optimizer.step
            running_loss += loss.item() * inputs.size(0)
            running_corrects += torch.sum(preds == labels.data)

        epoch_loss = running_loss / len(dataloader.dataset)
        epoch_acc = running_corrects.double() / len(dataloader.dataset)

        print('Epoch {}/{} Loss: {:.4f} Acc: {:.4f}'.format(epoch+1, num_epochs, epoch_loss, epoch_acc))

    return model

当初咱们能够开始训练模型:

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

model = train_model(model, trainloader, criterion, optimizer, num_epochs=25)

6. 评估模型

训练实现后,咱们须要评估模型的性能。咱们将应用测试数据集进行评估:

correct = 0
total = 0

with torch.no_grad():
    for inputs, labels in testloader:
        outputs = model(inputs)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print('Accuracy of the model on the test images: %d %%' % (100 * correct / total))

迁徙学习使咱们可能利用预训练模型的常识,从而更快地训练新工作。在本文中,咱们应用了 PyTorch 库和预训练的 ResNet-18 模型对 CIFAR-10 数据集进行分类。这种办法能够节俭大量计算资源和工夫,同时还能取得很好的性能。

退出移动版