PyTorch ResNet-50 模型训练代码详解 - GPU加速与自定义分类层

这段代码是一个经典的模型训练代码，用于在 PyTorch 中训练一个基于 ResNet-50 的分类模型。\n\n1. DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu"): 这行代码用于设置设备，如果有可用的 CUDA 设备，则使用 GPU，否则使用 CPU。\n\n2. model = models.resnet50(pretrained=True): 这行代码用于加载预训练的 ResNet-50 模型。\n\n3. in_features = model.fc.in_features: 这行代码用于获取 ResNet-50 模型最后一层全连接层的输入特征数。\n\n4. model.fc = nn.Sequential(nn.Linear(in_features, 256), nn.ReLU(), nn.Linear(256, 4), nn.LogSoftmax(dim=1)): 这行代码用于替换 ResNet-50 模型的最后一层全连接层，将输入特征数改为 256，输出类别数改为 4，并添加 ReLU 激活函数和 LogSoftmax 层。\n\n5. model = model.to(DEVICE): 这行代码将模型迁移到指定的设备，即将模型加载到 GPU（如果可用）或 CPU 上。\n\n6. loss_fn = nn.CrossEntropyLoss(): 这行代码定义了损失函数，使用交叉熵损失函数。\n\n7. loss_fn = loss_fn.to(DEVICE): 这行代码将损失函数迁移到指定的设备，即将损失函数加载到 GPU（如果可用）或 CPU 上。\n\n8. optimizer = optim.Adam(model.fc.parameters(), lr=0.003): 这行代码定义了优化器，使用 Adam 优化器，学习率为 0.003，只优化模型最后一层全连接层的参数。\n\n最后，我们可以使用这些代码来训练模型，具体的训练步骤包括前向传播、计算损失、反向传播和更新模型参数。