Vision Transformer 猫狗

### 实现猫狗分类的 Vision Transformer 模型 Vision Transformer（ViT）是一种基于自注意力机制的架构，用于处理图像数据。相比于传统的卷积神经网络（CNN），ViT 提供了一种全新的视角来捕捉全局特征[^1]。 #### 数据准备 为了构建一个 ViT 模型来进行猫狗分类任务，首先需要准备好相应的数据集。常用的数据集包括 Kaggle 上的 **Cat vs Dog Classification Dataset** 或者其他类似的公开数据源。这些数据通常由标记好的猫和狗图片组成。 #### 构建模型 以下是使用 PyTorch 来实现一个简单的 ViT 模型进行二分类的任务： ```python import torch from torch import nn, optim from torchvision import transforms, datasets from timm.models.vision_transformer import vit_base_patch16_224 # 使用预定义的 ViT 基础结构 # 定义超参数 batch_size = 32 image_size = 224 num_classes = 2 # 猫和狗两类 # 加载预训练的 ViT 模型 model = vit_base_patch16_224(pretrained=True) # 替换最后的全连接层以适应我们的分类任务 model.head = nn.Linear(model.head.in_features, num_classes) # 设置设备 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model.to(device) # 定义损失函数和优化器 criterion = nn.CrossEntropyLoss() optimizer = optim.AdamW(model.parameters(), lr=0.001) # 图像变换操作 data_transforms = { 'train': transforms.Compose([ transforms.RandomResizedCrop(image_size), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]), 'val': transforms.Compose([ transforms.Resize((image_size, image_size)), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) } # 加载数据集 dataset_path = './cat_dog_dataset' datasets_dict = {x: datasets.ImageFolder(f'{dataset_path}/{x}', data_transforms[x]) for x in ['train', 'val']} dataloaders_dict = {x: torch.utils.data.DataLoader(datasets_dict[x], batch_size=batch_size, shuffle=True) for x in ['train', 'val']} # 训练过程 def train_model(model, dataloaders, criterion, optimizer, device, num_epochs=10): model.train() best_acc = 0.0 for epoch in range(num_epochs): running_loss = 0.0 correct_preds = 0 for inputs, labels in dataloaders['train']: inputs, labels = inputs.to(device), labels.to(device) optimizer.zero_grad() outputs = model(inputs) loss = criterion(outputs, labels) _, preds = torch.max(outputs, 1) loss.backward() optimizer.step() running_loss += loss.item() * inputs.size(0) correct_preds += torch.sum(preds == labels.data) epoch_loss = running_loss / len(dataloaders['train'].dataset) epoch_acc = correct_preds.double() / len(dataloaders['train'].dataset) print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {epoch_loss:.4f}, Accuracy: {epoch_acc:.4f}') return model # 开始训练 trained_model = train_model(model, dataloaders_dict, criterion, optimizer, device=device, num_epochs=10) ``` 上述代码展示了如何加载预训练的 ViT 模型，并将其调整为适合特定任务的形式。注意，在实际应用中可能还需要进一步微调学习率、批量大小以及其他超参数。 #### 性能提升建议 由于原始 ViT 的计算复杂度较高[^3]，可以考虑一些改进措施： - 减少输入分辨率。 - 应用混合专家策略减少冗余计算。 - 利用更高效的变体如 Swin Transformer 或 DeiT。 ---