用pytorch 搭建一个coco数据集的训练脚本

以下是一个基于 PyTorch 框架搭建的 COCO 数据集训练脚本的示例代码：

import torch
import torchvision.models as models
import torchvision.transforms as transforms
import torchvision.datasets as datasets
from torch.utils.data import DataLoader

# 定义数据预处理方式
transform = transforms.Compose([
    transforms.Resize(256),
    transforms.RandomCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

# 加载 COCO 数据集
coco_train = datasets.CocoDetection(root='/path/to/coco', annFile='/path/to/annotations/train.json', transform=transform)
coco_val = datasets.CocoDetection(root='/path/to/coco', annFile='/path/to/annotations/val.json', transform=transform)

# 定义训练和验证数据的 DataLoader
train_loader = DataLoader(coco_train, batch_size=32, shuffle=True)
val_loader = DataLoader(coco_val, batch_size=32, shuffle=False)

# 定义模型
model = models.resnet50(pretrained=True)
num_classes = len(coco_train.coco.cats)
model.fc = torch.nn.Linear(model.fc.in_features, num_classes)

# 定义损失函数和优化器
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

# 训练模型
for epoch in range(10):
    for i, (images, labels) in enumerate(train_loader):
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'.format(epoch+1, 10, i+1, len(train_loader), loss.item()))

# 在验证集上测试模型
with torch.no_grad():
    correct = 0
    total = 0
    for images, labels in val_loader:
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()
    print('Accuracy of the network on the {} validation images: {} %'.format(len(val_loader), 100 * correct / total))

在这个示例代码中，我们首先定义了数据预处理方式，然后使用 `datasets.CocoDetection` 加载 COCO 数据集，并使用 `DataLoader` 定义训练和验证数据的 DataLoader。我们使用了 `models.resnet50` 作为模型，并将最后一层的输出改为与 COCO 数据集中的类别数相同。我们使用了交叉熵损失函数和随机梯度下降优化器进行训练。在每个 epoch 中，我们通过遍历训练 DataLoader 中的数据来更新模型参数，并打印出当前 epoch 和训练步骤的损失值。在训练完成后，我们在验证集上测试模型，并输出模型在验证集上的准确率。