写出coco数据集的图像分类训练函数

我们可以使用 PyTorch 中的 torchvision 库来训练 COCO 数据集上的图像分类模型。

下面是一个示例训练函数：

import torch
import torchvision
from torchvision.models import resnet50

def train_coco_image_classifier(train_dataset, val_dataset, batch_size, num_epochs):
    # 创建模型
    model = resnet50(pretrained=True)
    # 将模型的最后一层的输出维度设置为 COCO 数据集的类别数
    model.fc = torch.nn.Linear(model.fc.in_features, 80)
    # 定义损失函数和优化器
    criterion = torch.nn.CrossEntropyLoss()
    optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
    # 将模型移动到 GPU 上
    model = model.to(device)
    for epoch in range(num_epochs):
        for i, (inputs, labels) in enumerate(train_dataloader):
            # 将输入和标签移动到 GPU 上
            inputs = inputs.to(device)
            labels = labels.to(device)
            # 清空梯度
            optimizer.zero_grad()
            # 前向传播
            outputs = model(inputs)
            # 计算损失
            loss = criterion(outputs, labels)
            # 反向传播
            loss.backward()
            # 更新参数
            optimizer.step()
            # 输出训练进度
            if i % 100 == 0:
                print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'
                      .format(epoch+1, num_epochs, i+1, total_step, loss.item()))
    # 在验证集上评估模型
    with torch.no_grad():
        correct = 0
        total = 0
        for inputs, labels in val_dataloader:
            inputs = inputs.to(device)
            labels = labels.to(device)
            outputs = model(inputs)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
        print