用pytorch编写一个动物识别的代码

在PyTorch中编写一个简单的动物识别代码，通常会涉及以下几个步骤：

1. **数据准备**：
首先需要一个包含各种动物图片的数据集，例如ImageNet。使用`torchvision.datasets`加载数据，并通过`DataLoader`对数据进行批次划分和预处理。

import torch
from torchvision import datasets, transforms

data_transforms = {
    'train': transforms.Compose([
        transforms.RandomResizedCrop(224),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
    'val': transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])
}

dataset = datasets.ImageFolder(root='./data', transform=data_transforms['train'])
dataloader = torch.utils.data.DataLoader(dataset, batch_size=32, shuffle=True)

2. **构建模型**：
使用`torch.nn.Module`作为基础，可以选择预训练的卷积神经网络模型如`torchvision.models.resnet18`或`vgg16`，并添加全连接层用于分类。

import torch.nn as nn
import torchvision.models as models

model = models.resnet18(pretrained=True)
num_classes = len(dataset.classes)  # 确定类别数量
model.fc = nn.Linear(model.fc.in_features, num_classes)

3. **训练模型**：
定义损失函数（如交叉熵Loss）和优化器（如SGD、Adam等），然后进行训练过程。

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

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

for epoch in range(num_epochs):
    for images, labels in dataloader:
        images, labels = images.to(device), labels.to(device)

        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

    print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item()}')

4. **评估模型**：
在验证集上评估模型性能，并保存最佳模型。

# 评估阶段
model.eval()
correct = 0
total = 0
with torch.no_grad():
    for data in dataloader_val:
        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()

accuracy = 100 * correct / total
print(f'Accuracy of the model on the validation set: {accuracy}%')