乐趣区

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

jiezi

在本文中，咱们将介绍如何应用 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 数据集进行分类。这种办法能够节俭大量计算资源和工夫，同时还能取得很好的性能。

退出移动版