【AI大模型】Transformers大模型库（十五）：timm库

LDG_AGI

发布于 2024-08-13 16:20:39

970

发布于 2024-08-13 16:20:39

一、引言

这里的Transformers指的是huggingface开发的大模型库，为huggingface上数以万计的预训练大模型提供预测、训练等服务。

🤗 Transformers 提供了数以千计的预训练模型，支持 100 多种语言的文本分类、信息抽取、问答、摘要、翻译、文本生成。它的宗旨是让最先进的 NLP 技术人人易用。 🤗 Transformers 提供了便于快速下载和使用的API，让你可以把预训练模型用在给定文本、在你的数据集上微调然后通过 model hub 与社区共享。同时，每个定义的 Python 模块均完全独立，方便修改和快速研究实验。 🤗 Transformers 支持三个最热门的深度学习库： Jax, PyTorch 以及 TensorFlow — 并与之无缝整合。你可以直接使用一个框架训练你的模型然后用另一个加载和推理。

本文重点介绍Hugging Face的timm库用法

二、timm库

2.1 概述

Hugging Face的timm库是一个用于计算机视觉的模型库，它提供了大量预训练的图像识别模型，以高效、易用为特点。

2.2 使用方法

2.2.1 安装timm库

首先，确保您已经安装了timm库。可以通过pip命令安装：

   pip install timm

2.2.2 导入必要的库

在Python脚本中，您需要导入timm库以及PyTorch库来构建和训练模型。

   import torch
   import timm
   from torch.utils.data import DataLoader
   from torchvision import datasets, transforms

2.2.3 数据预处理

准备数据集并进行预处理，例如缩放、归一化等。

   transform = transforms.Compose([
       transforms.Resize(256),
       transforms.CenterCrop(224),
       transforms.ToTensor(),
       transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
   ])

   dataset = datasets.ImageFolder('your_dataset_path', transform=transform)
   dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

2.2.4 选择模型

timm库提供了很多模型，例如EfficientNet，ResNet等，这里以EfficientNet为例。

   model = timm.create_model('efficientnet_b0', pretrained=True)
   model.classifier = torch.nn.Linear(model.classifier.in_features, num_classes)  # 修改分类层，num_classes为您的类别数

2.2.5 损失函数和优化器

   criterion = torch.nn.CrossEntropyLoss()
   optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

2.2.6 训练模型

定义训练循环:

   device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
   model.to(device)

   for epoch in range(num_epochs):  # num_epochs为训练轮数
       for inputs, labels in dataloader:
           inputs, labels = inputs.to(device), labels.to(device)
           optimizer.zero_grad()
           outputs = model(inputs)
           loss = criterion(outputs, labels)
           loss.backward()
           optimizer.step()
       print(f'Epoch {epoch+1}/{num_epochs}, Loss: {loss.item()}')