Deepfake-FFDI-图像赛题 baseline（3）加载模型 部分 @Datawhale AI 夏令营

import timm
model = timm.create_model('resnet18', pretrained=True, num_classes=2)
model = model.cuda(

 1. 创建数据加载器

train_loader = torch.utils.data.DataLoader(
    FFDIDataset(train_label['path'].head(1000), train_label['target'].head(1000), 
            transforms.Compose([
                        transforms.Resize((256, 256)),
                        transforms.RandomHorizontalFlip(),
                        transforms.RandomVerticalFlip(),
                        transforms.ToTensor(),
                        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])
    ), batch_size=40, shuffle=True, num_workers=4, pin_memory=True
)

val_loader = torch.utils.data.DataLoader(
    FFDIDataset(val_label['path'].head(1000), val_label['target'].head(1000), 
            transforms.Compose([
                        transforms.Resize((256, 256)),
                        transforms.ToTensor(),
                        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])
    ), batch_size=40, shuffle=False, num_workers=4, pin_memory=True
)

• train_loader:

  • 使用 FFDIDataset 创建训练集的数据加载器。
  • 图像经过预处理和数据增强，包括调整大小、随机水平翻转、随机垂直翻转、转换为张量以及归一化。
  • batch_size 设置为 40，数据会被打乱（shuffle=True），并使用 4 个工作线程加载数据（num_workers=4）。
  • pin_memory=True 使得数据加载更高效。

• val_loader:

  • 使用 FFDIDataset 创建验证集的数据加载器。
  • 图像只经过调整大小、转换为张量和归一化，没有数据增强。
  • batch_size 设置为 40，数据不会被打乱（shuffle=False），并使用 4 个工作线程加载数据（num_workers=4）。
  • pin_memory=True 使得数据加载更高效

2. 定义损失函数和优化器

criterion = nn.CrossEntropyLoss().cuda()
optimizer = torch.optim.Adam(model.parameters(), 0.005)

• criterion:

  • 使用交叉熵损失作为损失函数，并将其移动到 GPU 上（使用 .cuda()）。

• optimizer:

  • 使用 Adam 优化器，学习率设置为 0.005。

3. 定义学习率调度器

scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=4, gamma=0.85)

scheduler:

• 使用学习率调度器 StepLR，每 4 个周期将学习率乘以 0.85。

4. 训练和验证循环

best_acc = 0.0
for epoch in range(2):
    scheduler.step()
    print('Epoch: ', epoch)

    train(train_loader, model, criterion, optimizer, epoch)
    val_acc = validate(val_loader, model, criterion)
    
    if val_acc.avg.item() > best_acc:
        best_acc = round(val_acc.avg.item(), 2)
        torch.save(model.state_dict(), f'./model_{best_acc}.pt')

• 初始化最佳准确率:

  • best_acc 初始化为 0.0。

• 训练周期循环:

  • 进行 2 个训练周期（range(2)）。
  • 每个周期开始时，更新学习率调度器（scheduler.step()）。
  • 打印当前周期数。

• 训练和验证:

  • 使用 train 函数在训练集上训练模型。
  • 使用 validate 函数在验证集上评估模型，返回验证准确率 val_acc。

• 保存最佳模型:

  • 如果当前验证准确率超过了之前的最佳准确率，则更新 best_acc 并保存当前模型的状态字典到文件 ./model_{best_acc}.pt。

训练和验证函数

假设 train 和 validate 函数是之前已经定义好的，用于训练模型和在验证集上评估模型的性能。具体实现可能类似于之前回答中的 train 和 validate 函数。

下面这段代码展示了如何在训练和验证阶段之后，使用验证数据集（或测试数据集）来进行模型的预测，并将预测结果保存到一个 CSV 文件中。具体来说，这段代码的流程如下：

代码结构

1. 创建测试数据加载器
2. 进行预测并保存结果

1. 创建测试数据加载器

test_loader = torch.utils.data.DataLoader(
    FFDIDataset(val_label['path'], val_label['target'], 
            transforms.Compose([
                        transforms.Resize((256, 256)),
                        transforms.ToTensor(),
                        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])
    ), batch_size=40, shuffle=False, num_workers=4, pin_memory=True
)

• test_loader:
  • 使用 FFDIDataset 创建测试集（在这里实际上是验证集）的数据加载器。
  • 图像经过调整大小、转换为张量和归一化，没有数据增强。
  • batch_size 设置为 40，数据不会被打乱（shuffle=False），并使用 4 个工作线程加载数据（num_workers=4）。
  • pin_memory=True 使得数据加载更高效。

2. 进行预测并保存结果

val_label['y_pred'] = predict(test_loader, model, 1)[:, 1]
val_label[['img_name', 'y_pred']].to_csv('submit.csv', index=None)

• 预测:

  • 使用 predict 函数在测试数据加载器上进行预测。
  • predict 函数接受 test_loader、model 和一个附加参数（在这里是 1），返回预测结果。假设返回的是一个二维数组或张量，其中每一行是一个样本的预测结果。
  • 选择预测结果的第二列（[:, 1]），并将其赋值给 val_label['y_pred']。这可能是一个二分类问题，其中第二列代表正类的概率或分数。

• 保存结果:

  • 从 val_label 数据框中选择 img_name 和 y_pred 两列，并将其保存为 CSV 文件 submit.csv。
  • index=None 参数表示不保存行索引。

预测函数

假设 predict 函数已经定义好，用于在测试数据加载器上进行预测。一个可能的实现如下：

def predict(test_loader, model, device):
    model.eval()  # 将模型设为评估模式
    all_preds = []
    
    with torch.no_grad():  # 禁用梯度计算，节省内存和计算资源
        for images, _ in test_loader:
            images = images.to(device)  # 将图像移动到指定设备（如GPU）
            outputs = model(images)  # 前向传播
            preds = torch.softmax(outputs, dim=1).cpu().numpy()  # 获取预测结果并转换为numpy数组
            all_preds.append(preds)
    
    return np.concatenate(all_preds, axis=0)  # 拼接所有预测结果

# 设备可以是 'cuda' 或 'cpu'
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model.to(device)

到这里所有的baseline的详细注释就结束了。