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()
)。
- 使用交叉熵损失作为损失函数,并将其移动到 GPU 上(使用
-
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()
)。 - 打印当前周期数。
- 进行 2 个训练周期(
-
训练和验证:
- 使用
train
函数在训练集上训练模型。 - 使用
validate
函数在验证集上评估模型,返回验证准确率val_acc
。
- 使用
-
保存最佳模型:
- 如果当前验证准确率超过了之前的最佳准确率,则更新
best_acc
并保存当前模型的状态字典到文件./model_{best_acc}.pt
。
- 如果当前验证准确率超过了之前的最佳准确率,则更新
训练和验证函数
假设 train
和 validate
函数是之前已经定义好的,用于训练模型和在验证集上评估模型的性能。具体实现可能类似于之前回答中的 train
和 validate
函数。
下面这段代码展示了如何在训练和验证阶段之后,使用验证数据集(或测试数据集)来进行模型的预测,并将预测结果保存到一个 CSV 文件中。具体来说,这段代码的流程如下:
代码结构
- 创建测试数据加载器
- 进行预测并保存结果
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的详细注释就结束了。