当前位置：首页 > article >正文

timm使用笔记

article 2024/11/2 20:49:21

timm（Timm is a model repository for PyTorch）是一个 PyTorch 原生实现的计算机视觉模型库。它提供了预训练模型和各种网络组件，可以用于各种计算机视觉任务，例如图像分类、物体检测、语义分割等等。timm（库提供了预训练模型、模型构建块和模型训练的实用工具。timm库可以帮助开发者快速构建和训练深度学习模型，同时支持多种图像分类、分割和检测任务，特别是结合torch和torchvision的使用，对你训练模型，事半功倍。

timm 的特点如下：

PyTorch 原生实现：timm 的实现方式与 PyTorch 高度契合，开发者可以方便地使用 PyTorch 的 API 进行模型训练和部署。
轻量级的设计：timm 的设计以轻量化为基础，根据不同的计算机视觉任务，提供了多种轻量级的网络结构。
大量的预训练模型：timm 提供了大量的预训练模型，可以直接用于各种计算机视觉任务。
多种模型组件：timm 提供了各种模型组件，如注意力模块、正则化模块、激活函数等等，这些模块都可以方便地插入到自己的模型中。
高效的代码实现：timm 的代码实现高效并且易于使用。

需要注意的是，timm 是一个社区驱动的项目，它由计算机视觉领域的专家共同开发和维护。在使用时需要遵循相关的使用协议。

1. 安装

pip install timm

2. 准备数据集

首先，我们需要准备CIFAR-10数据集。我们可以使用torchvision库来下载和加载数据集：

import torch
import torchvision
import torchvision.transforms as transforms

# 数据预处理

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))
])

# 加载CIFAR-10数据集

trainset = torchvision.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 = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=100, shuffle=False, num_workers=2)

3. 加载预训练模型

timm库提供了多个预训练模型，这些模型可以在ImageNet等数据集上进行预训练，也可以在其他数据集上进行微调。

加载预训练模型的代码非常简单，下面我们加载需要的预训练模型权重：

import timm

m = timm.create_model('vgg16', pretrained=True)
m.eval()

上面代码就会创建一个VGG-16的预训练模型。

3. 训练模型

现在我们可以开始训练模型。我们将使用交叉熵损失函数和Adam优化器：

import torch.optim as optim


criterion = torch.nn.CrossEntropyLoss()

optimizer = optim.Adam(model.parameters(), lr=0.001)

# 训练模型
num_epochs = 10

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model.to(device)

for epoch in range(num_epochs):
    running_loss = 0.0

    for i, data in enumerate(trainloader, 0):
        inputs, labels = data
        inputs, labels = inputs.to(device), labels.to(device)
        optimizer.zero_grad()
        outputs = model(inputs)

        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()

    print(f"Epoch {epoch + 1}, Loss: {running_loss / (i + 1)}")

4. 测试模型

训练完成后，我们可以使用测试数据集评估模型的性能：

correct = 0
total = 0

model.eval()

with torch.no_grad():

    for data in testloader:
        images, labels = data
        images, labels = images.to(device), labels.to(device)

        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)

        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print(f"Accuracy on test set: {100 * correct / total}%")