from model import AttU_Net as UNet from utils.dataset import ISBI_Loader from torch import optim import torch.nn as nn import torch def train_net(net, device, data_path, epochs=40, batch_size=1, lr=0.00001): # 加载训练集 isbi_dataset = ISBI_Loader(data_path) train_loader = torch.utils.data.DataLoader(dataset=isbi_dataset, batch_size=batch_size, shuffle=True) # 定义RMSprop算法 optimizer = optim.RMSprop(net.parameters(), lr=lr, weight_decay=1e-8, momentum=0.9) # 定义了一个名为`optimizer`的优化器,它是基于RMSprop算法的。这个优化器用于在训练神经网络时更新网络的权重。具体参数如下: # `net.parameters()`:这是神经网络`net`的所有可学习参数,包括权重和偏置。 # `lr`:学习率,用于控制每次更新权重时的步长。 # `weight_decay`:权重衰减,用于控制权重更新的幅度。 # `momentum`:动量,用于加速梯度下降过程。 # RMSProp 是一种自适应的优化算法 # RMSprop 主要通过对权重更新公式进行修改,引入了指数加权平均值来平滑梯度,w = w - learning_rate * sqrt(cache + epsilon) * grad, # 其中,`learning_rate` 是学习率,`cache` 是梯度的一阶矩(均值),`grad` 是当前梯度,`epsilon` 是一个很小的常数,用于避免除以零的情况。 # RMSprop 的优点在于,它能够适应不同梯度的大小,当梯度较大时,更新幅度较小,避免权重更新过大;当梯度较小时,更新幅度较大,使权重收敛更快。 # 定义Loss算法 criterion = nn.BCEWithLogitsLoss() # 这个损失函数通常用于二分类任务中,这个项目只需要分割出细胞边缘,也就是一个很简单的二分类任务 # best_loss统计,初始化为正无穷 best_loss = float('inf') # 训练epochs次 for epoch in range(epochs): # 训练模式 net.train() # 按照batch_size开始训练 for image, label in train_loader: optimizer.zero_grad() # 将模型参数的梯度归零 # 将数据拷贝到device中 image = image.to(device=device, dtype=torch.float32) # 将输入图像(image)和标签(label)从CPU内存复制到GPU显存中,并转换为浮点类型(torch.float32)。 label = label.to(device=device, dtype=torch.float32) # 使用网络参数,输出预测结果 pred = net(image) # 计算loss loss = criterion(pred, label) print('Loss/train', loss.item()) # 保存loss值最小的网络参数 if loss < best_loss: best_loss = loss torch.save(net.state_dict(), 'best_model.pth') # 更新参数 loss.backward() optimizer.step() # 参数更新 if __name__ == "__main__": # 选择设备,有cuda用cuda,没有就用cpu device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # 加载网络,图片单通道1,分类为1。 net = UNet(img_ch=1, output_ch=1) # 将网络拷贝到deivce中 net.to(device=device) # 指定训练集地址,开始训练 data_path = "D:\\AAAmodel\\model\\data\\train" train_net(net, device, data_path)
在提供的代码中,定义了一个名为 train_net
的函数,其目的是训练一个基于PyTorch的UNet模型,用于图像分割任务。以下是对代码的详细解释:
-
加载数据集:首先,代码通过
ISBI_Loader
类加载训练集数据,该类可能定义了如何加载和处理图像数据集。然后,使用torch.utils.data.DataLoader
创建了一个数据加载器train_loader
,它负责以指定的批量大小batch_size
加载数据,并在每个epoch开始时对数据进行打乱shuffle=True
。 -
定义优化器:使用
torch.optim.RMSprop
定义了一个优化器optimizer
。RMSprop 是一种自适应学习率优化算法,适用于非凸优化问题。优化器的参数包括:net.parameters()
:模型的所有可学习参数。lr=lr
:学习率,控制权重更新的步长。weight_decay=1e-8
:权重衰减,用于正则化以防止过拟合。momentum=0.9
:动量,用于加速梯度下降过程,提高收敛速度。
-
定义损失函数:使用
torch.nn.BCEWithLogitsLoss
定义了二元分类的损失函数criterion
。这个损失函数结合了sigmoid
激活函数和binary cross entropy
损失,适用于二分类问题。 -
训练过程:代码通过一个外部循环遍历指定的轮数
epochs
。在每个epoch中,模型被设置为训练模式net.train()
,然后通过内部循环遍历数据加载器中的所有批次。对于每个批次:- 使用
optimizer.zero_grad()
清空梯度,为反向传播做准备。 - 将图像和标签数据移动到指定的设备(GPU或CPU)上。
- 通过模型进行前向传播,得到预测结果
pred
。 - 计算预测结果和真实标签之间的损失
loss
。 - 如果当前损失小于之前记录的最小损失
best_loss
,则保存当前模型的参数到文件best_model.pth
。 - 执行反向传播
loss.backward()
并更新模型参数optimizer.step()
。
- 使用
-
保存最佳模型:在训练过程中,如果模型在当前epoch的损失低于之前记录的最小损失,则将当前模型的参数保存为
best_model.pth
。这有助于在训练结束后恢复最佳性能的模型。 -
主函数:在
if __name__ == "__main__":
块中,代码首先确定使用GPU还是CPU进行训练。然后,实例化UNet模型net
,并将模型移动到选定的设备上。最后,指定训练集的路径data_path
并调用train_net
函数开始训练过程。
整体而言,这段代码展示了如何使用PyTorch框架和UNet模型进行图像分割任务的训练。代码中包含了数据加载、模型定义、优化器和损失函数的选择、训练过程以及模型保存等关键步骤。12345678