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

timm使用笔记

news 2025/7/4 10:09:40

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 timmm = timm.create_model('vgg16', pretrained=True)
m.eval()

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

3. 训练模型

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

import torch.optim as optimcriterion = torch.nn.CrossEntropyLoss()optimizer = optim.Adam(model.parameters(), lr=0.001)# 训练模型
num_epochs = 10device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model.to(device)for epoch in range(num_epochs):running_loss = 0.0for i, data in enumerate(trainloader, 0):inputs, labels = datainputs, 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 = 0model.eval()with torch.no_grad():for data in testloader:images, labels = dataimages, 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}%")