参数设置
if __name__ == "__main__":
基本参数
Cuda = True
distributed = False # 用于指定是否使用单机多卡分布式运行
sync_bn = False # 是否使用sync_bn，DDP模式多卡可用
distributed用于指定是否使用单机多卡分布式运行，终端指令仅支持Ubuntu。CUDA_VISIBLE_DEVICES用于在Ubuntu下指定显卡。
Windows系统下默认使用DP模式调用所有显卡，不支持DDP。
- DP模式：设置
distributed = False在终端中输入 CUDA_VISIBLE_DEVICES=0,1 python train.py- DDP模式：设置
distributed = True在终端中输入 CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.launch --nproc_per_node=2 train.py
fp16 = False # 是否使用混合精度训练
fp16混合精度指的是使用半精度浮点数（half-precision floating-point）来进行计算的一种技术。浮点数是一种表示实数的数据类型，用于存储和计算带有小数点的数值。在深度学习和大规模数据处理等计算密集型任务中，深层神经网络模型和大规模矩阵计算会涉及大量的浮点数运算。
传统上，这些计算都是使用单精度浮点数（32位浮点数）进行的，而半精度浮点数可以显著减少计算过程中的内存和计算开销使用fp16混合精度的技术意味着将神经网络中的一些参数和激活值转换为半精度浮点数，并在计算过程中使用这些半精度值。这样做可以大幅减少数据在内存和缓存中的存储需求，同时加快计算速度，因为半精度浮点数在现代GPU硬件上的计算速度通常比单精度浮点数更快。
然而，由于半精度浮点数的表示范围较小，可能会导致计算中的数值下溢或上溢。因此，在应用fp16混合精度时，需要小心处理数值的范围，并可能需要使用特殊的技巧来缓解这些问题，如使用缩放因子来保持数值范围合适。
num_classes = 2 # 目标分类个数+1
phi = "b0" # 所使用的的主干网络：b0、b1、b2、b3、b4、b5
预训练
pretrained = False # 是否使用主干网络的预训练权重（是否使用预训练）
如果设置了
model_path，则主干的权值无需加载，pretrained的值无意义。
如果不设置model_path，pretrained = True，此时仅加载主干开始训练。
如果不设置model_path，pretrained = False，Freeze_Train = Fasle，此时从0开始训练，且没有冻结主干的过程。
"pretrained"预训练是一种常用的技术，它指的是在一个较大的数据集上预先训练好的神经网络模型。以学习一些普遍的特征和知识。
backbone network主干网络是指构成整个神经网络模型核心部分的部分。例如，在图像分类任务中，主干
网络通常是用来提取图像特征的卷积神经网络（CNN）部分；在自然语言处理任务中，主干网络可能是用来处理文本序列的循环神经网络（RNN）或者注意力机制（Transformer）部分。
模型使用了主干网络的预训练权重时，意味着在模型构建的过程中，这个主干网络部分的权重被加载自预训练模型中。这样做的好处是，由于预训练模型已经在大规模数据上学习到了一些通用的特征和知识，将这些预训练权重用于当前任务的模型构建可以为模型提供更好的初始状态，加速模型在目标任务上的训练过程，并且可能提高模型的性能。
- 模型的 预训练权重 对不同数据集是通用的，因为特征是通用的
- 预训练权重对于99%的情况都必须要用，不用的话主干部分的权值太过随机，特征提取效果不明显，网络训练的结果也不会好
- 训练自己的数据集时提示维度不匹配正常，预测的东西都不一样了自然维度不匹配
- 如果训练过程中存在中断训练的操作，可以将
model_path设置成logs文件夹下的权值文件，将已经训练了一部分的权值再次载入。- 同时修改下方的 冻结阶段 或者 解冻阶段 的参数，来保证模型epoch的连续性。 当
model_path = ' '的时候不加载整个模型的权值。 此处使用的是整个模型的权重，因此是在train.py进行加载的，pretrain不影响此处的权值加载。- 如果想要让模型从主干的预训练权值开始训练，则设置
model_path = ' '，pretrain = True，此时仅加载主干。- 如果想要让模型从0开始训练，则设置
model_path = ' '，pretrain = Fasle，Freeze_Train =Fasle，此时从0开始训练，且没有冻结主干的过程。- 一般来讲，网络从0开始的训练效果会很差，因为权值太过随机，特征提取效果不明显，因此非常、非常、非常不建议大家从0开始训练！
- 如果一定要从0开始，可以了解imagenet数据集，首先训练分类模型，获得网络的主干部分权值，分类模型的主干部分和该模型通用，基于此进行训练。
model_path = "model_data/segformer_b0_weights_voc.pth" # 加载的权值文件路径
input_shape = [512, 512] # 输入图片的大小
冻结和解冻训练
在训练语义分割模型的过程中，冻结阶段和解冻阶段是为了逐步训练模型的不同部分，以满足机器性能有限的情况下进行训练。
Freeze_Train = True # 是否进行冻结训练， 默认先冻结主干训练后解冻训练。
1. 冻结阶段：
在冻结阶段，将模型的一部分或全部层设置为不可训练（冻结状态）。通常情况下，我们会将模型的主干特征提取网络（Backbone）部分冻结，只训练语义分割头部（Segmentation Head）部分。主干网络通常是一个预训练的庞大模型（如ResNet、EfficientNet等），它包含大量参数，需要较长的时间来训练。
冻结主干网络意味着在这个阶段不更新主干网络的参数，只更新语义分割头部的参数，以降低训练的计算成本和时间。
在冻结阶段，模型的反向传播只发生在语义分割头部，而不影响主干网络的参数因此，训练速度会较快，但由于冻结的主干网络可能不够灵活，可能会限制模型在训练数据上的拟合能力。
2. 解冻阶段：
在解冻阶段，解除对主干网络的冻结，使其所有参数都可以进行训练这个阶段的目标是允许主干网络的参数根据当前任务进行微调，以更好地适应特定的分割任务和训练数据。
在解冻阶段，整个模型的所有参数都会被更新，包括主干网络和语义分割头部的参数。这样做的好处是，模型能够更好地适应训练数据，提高模型的拟合能力和表现。
网络运行机制：
在整个训练过程中，网络的运行机制会根据不同阶段进行调整：
-
冻结阶段：只有语义分割头部的参数会根据训练数据进行更新，主干网络的参数保持不变。这样，训练过程中只计算语义分割头部的梯度，大部分计算量集中在这一部分，因此训练速度会较快。
-
解冻阶段：整个模型的所有参数都会根据训练数据进行更新，包括主干网络和语义分割头部。训练过程中将计算整个模型的梯度，这会增加计算量，训练速度可能会减慢。但由于主干网络此时可以根据任务微调，模型的表现有望得到进一步提升。
通过这种两阶段训练的方式，可以在机器性能有限的情况下，首先快速训练模型的头部部分，然后再逐步解冻主干网络，让整个模型得到更好的训练和拟合，同时兼顾训练效率和模型性能的平衡。
| 阶段 | 参数 | 解读 |
|---|---|---|
| 开始 | Init_Epoch | 模型当前开始的训练Epoch |
| 冻结 | Freeze_Epoch | 模型冻结训练的Epoch |
| Freeze_batch_size | 模型冻结训练的batch_size | |
| 解冻 | UnFreeze_Epoch | 模型总共训练的Epoch |
| Unfreeze_batch_size | 模型在解冻后的batch_size |
在深度学习中，“世代”（Epoch）和 “批量大小”（Batch Size）是两个重要的概念，它们涉及到训练数据的处理和模型参数的更新。
-
世代（Epoch）：
一个世代表示将整个训练数据集中的所有样本都用于训练一次。在训练神经网络时，通常会将数据集分成若干批次（batches），每个批次包含一定数量的样本。一个世代由多个批次组成，每个批次的样本都会用于更新模型的参数。在训练过程中，多个世代的迭代可以使模型更好地适应训练数据，提高模型的泛化能力。
例如，假设有1000个样本的数据集，批量大小设置为100。那么一个世代将包含10个批次，每个批次包含100个样本，整个数据集中的样本都会在一个世代内使用。 -
批量大小（Batch Size）：
批量大小是指每个批次中包含的样本数量。它表示在每次模型的前向传播和反向传播过程中，同时处理的样本数目。较大的批量大小可以充分利用硬件并行性，加快训练速度，但可能需要更多的内存。较小的批量大小可以减少内存需求，但可能导致训练速度较慢。
举例来说，假设我们有一个包含1000个样本的数据集，希望使用批量大小为32进行训练。那么整个数据集将被分成32个批次，每个批次包含32个样本。在训练过程中，模型将依次处理每个批次，计算损失并更新参数，直到遍历完所有批次，完成一个周期的训练。
“批次”是一个泛指，表示一组同时处理的样本，而“批量大小”则是指具体的每个批次中包含的样本数量。在深度学习中，它们通常指代不同的概念
- 冻结训练需要的显存较小，显卡非常差的情况下，可设置Freeze_Epoch等于UnFreeze_Epoch，此时仅仅进行冻结训练。
- 从整个模型的预训练权重开始训练：
Adam：
（冻结）Init_Epoch = 0，Freeze_Epoch = 50，UnFreeze_Epoch = 100，Freeze_Train = True，optimizer_type = 'adam'，Init\_lr = 1e-4，weight_decay = 0。
（不冻结）Init_Epoch =0，UnFreeze_Epoch = 100，Freeze_Train = False，optimizer_type = 'adam'，Init_lr = 1e-4，weight_decay = 0。
AdamW：
（冻结）Init_Epoch =0，Freeze_Epoch = 50，UnFreeze_Epoch = 100，Freeze_Train =True，optimizer_type = 'adamw'，Init_lr = 1e-4，weight_decay =1e-2。
（不冻结）Init_Epoch = 0，UnFreeze_Epoch = 100，Freeze_Train = False，optimizer_type = 'adamw'，Init_lr = 1e-4，weight_decay = 1e-2。
UnFreeze_Epoch可以在100-300之间调整。- 从主干网络的预训练权重开始训练：
Adam：
（冻结）Init_Epoch = 0，Freeze_Epoch = 50，UnFreeze_Epoch = 100，Freeze_Train = True，optimizer_type = 'adam'，Init_lr = 1e-4，weight_decay = 0。
（不冻结）Init_Epoch = 0，UnFreeze_Epoch = 100，Freeze_Train =False，optimizer_type = 'adam'，Init_lr = 1e-4，weight_decay = 0。
AdamW：
（冻结）Init_Epoch = 0，Freeze_Epoch = 50，UnFreeze_Epoch =100，Freeze_Train = True，optimizer_type = 'adamw'，Init_lr = 1e-4，weight_decay = 1e-2。
（不冻结）Init_Epoch = 0，UnFreeze_Epoch =100，Freeze_Train = False，optimizer_type = 'adamw'，Init_lr =1e-4，weight_decay = 1e-2。
其中：由于从主干网络的预训练权重开始训练，主干的权值不一定适合语义分割，需要更多的训练跳出局部最优解。
UnFreeze_Epoch可以在100-300之间调整。
冻结阶段训练参数
Init_Epoch = 0 # 模型当前开始的训练世代，其值可以大于Freeze_Epoch，如设置：Init_Epoch = 60、Freeze_Epoch = 50、UnFreeze_Epoch = 100，会跳过冻结阶段，直接从60代开始，并调整对应的学习率。（断点续练时使用）
Freeze_Epoch = 50 # 模型冻结训练的Freeze_Epoch(当Freeze_Train=False时失效)
Freeze_batch_size = 16 # 模型冻结训练的batch_size(当Freeze_Train=False时失效)
此时模型的主干被冻结了，特征提取网络不发生改变，占用的显存较小，仅对网络进行微调
batch_size的设置：
在显卡能够接受的范围内，以大为好。显存不足与数据集大小无关，提示显存不足（OOM或者CUDA out of memory）请调小batch_size。
受到BatchNorm层影响，batch_size最小为2，不能为1。
正常情况下Freeze_batch_size建议为Unfreeze_batch_size的1-2倍。不建议设置的差距过大，因为关系到学习率的自动调整
解冻阶段训练参数
UnFreeze_Epoch = 100 # 模型总共训练的epoch
Unfreeze_batch_size = 8 # 模型在解冻后的batch_size
此时模型的主干不被冻结了，特征提取网络会发生改变，占用的显存较大，网络所有的参数都会发生改变
其它训练参数：学习率、优化器、学习率下降有关
Init_lr = 1e-4 # 模型的最大学习率
Min_lr = Init_lr * 0.01 # 模型的最小学习率
Init_lr
当使用Adam优化器时建议设置Init_lr=1e-4
当使用AdamW优化器时建议设置Init_lr=1e-4
Transformer系列不建议使用SGD
Min_lr
模型的最小学习率，默认为最大学习率的0.01
optimizer_type = "adamw" # 使用到的优化器种类，可选的有adam、adamw、sgd
momentum = 0.9 # 优化器内部使用到的momentum参数
weight_decay = 1e-2 # 权值衰减，可防止过拟合，adam会导致weight_decay错误，使用adam时建议设置为0
momentum
- 在深度学习中，momentum（动量）是一种优化算法的一部分，用于加速梯度下降的过程，特别是在训练神经网络时。
- 在梯度下降优化算法中，每次更新模型参数时，都会根据当前的梯度方向来进行参数更新。而动量优化算法引入了一个额外的因子（即动量值）来模拟物体在惯性作用下的运动。它的主要目的是解决梯度下降算法在梯度方向改变时可能出现震荡或缓慢收敛的问题。
- 具体来说，动量优化算法会在每次参数更新时，考虑之前的梯度方向，并将当前的梯度方向与之前的梯度方向进行结合。这样做的效果是，如果在连续几个迭代中梯度方向保持一致，那么更新的步伐就会加速；而如果在不同的迭代中梯度方向发生变化，更新的步伐就会减缓，从而在参数更新过程中形成更加平滑和稳定的路径。
- 通过引入动量，梯度下降的更新方向更加稳定，有助于加快收敛速度，并且可以在一定程度上避免陷入局部最优点。因此，
momentum是深度学习优化算法中常用的一种技术，帮助优化算法更快地找到全局最优点或者更好的解。
weight_decay
- 在深度学习中，weight decay（权重衰减）是一种用于正则化模型的技术，旨在减小模型的过拟合风险。
- 权重衰减通过在损失函数中添加一个正则化项，对模型的权重进行约束。这个正则化项是权重的平方和（L2范数）与一个较小的正则化系数相乘，即 weight_decay * ||w||2。其中，w表示模型的权重参数，||w||2表示权重的L2范数（即权重的平方和），weight_decay是一个用户定义的较小正数。
- 在训练模型时，优化器会通过最小化损失函数来调整模型的权重，这样模型可以更好地拟合训练数据。但是，如果模型的权重过大，可能会导致模型过于复杂，容易出现过拟合现象，即在训练数据上表现很好，但在未见过的数据上表现较差。
- 通过引入权重衰减，优化器在更新权重时会减少权重的大小，从而使模型的复杂度减小，有助于缓解过拟合问题。正则化项的大小由weight_decay控制，较小的weight_decay值意味着对权重的约束较弱，而较大的weight_decay值会更严格地限制权重。
- 需要注意的是，权重衰减在优化器的更新规则中等效于在损失函数中添加一个关于权重的惩罚项。实际上，权重衰减就是L2正则化的一种表现形式。
- 举例来说，如果
weight_decay=1e-2，那么权重衰减项就是模型中所有权重的平方和乘以1e-2，用来在训练过程中对模型的权重进行惩罚，以减少过拟合的风险。
lr_decay_type = 'cos' # 使用到的学习率下降方式，可选的有'step'、'cos'
lr_decay_type = 'cos'表示学习率（Learning Rate）的衰减方式是余弦退火（Cosine Annealing）。- 在深度学习的训练过程中，学习率是一个非常重要的超参数。学习率控制了参数在每次迭代（或每个批次）中更新的步长。较大的学习率可以加快训练速度，但可能导致不稳定的训练过程甚至发散；而较小的学习率可以稳定训练，但收敛速度可能较慢。
- 余弦退火是一种学习率调度策略，通过在训练过程中逐渐减小学习率，从而在训练初期使用较大的学习率进行快速收敛，然后逐渐减小学习率以更加稳定地优化模型。这种方式能够在训练后期使得模型更加精确。
- 具体来说，余弦退火将学习率进行周期性的余弦函数调度。学习率会在一个周期内逐渐减小，然后重新增大，在下一个周期内继续逐渐减小。通常情况下，这个周期的长度是训练的总世代数。
- 余弦退火的优势在于它能够在不增加训练时间的情况下提高模型的精确度。它通过在训练过程中适应性地调整学习率，避免了一开始就选择较小的固定学习率可能导致的训练停滞问题。
save_period = 5 # 多少个epoch保存一次权值
save_dir = 'logs' # 权值与日志文件保存的文件夹
eval_flag = True # 是否在训练时进行评估，评估对象为验证集,此处获得的mAP为验证集的mAP,与get_miou.py不同
eval_period = 5 # 代表多少个epoch评估一次，不建议频繁的评估
VOCdevkit_path = 'VOCdevkit' # VOCdevkit_path 数据集路径
dice_loss = False
dice_loss 是一个布尔值，用于控制是否使用 Dice Loss 来作为训练过程中的损失函数。
Dice Loss 是一种常用于图像分割任务的损失函数，它衡量预测结果与真实标签之间的相似度。Dice Loss 的计算方式是基于 Dice 系数（也称为 F1-Score），计算公式如下：
Dice 系数 = (2 * 预测结果与真实标签的交集大小) / (预测结果的大小 + 真实标签的大小)
Dice Loss 则是将 Dice 系数的值转换为一个损失值，一般来说，Dice Loss 的取值范围是 0 到 1，数值越小表示预测结果与真实标签之间的差异越大，数值越大表示相似度越高，即损失越小。
在代码中，如果 dice_loss 设置为 False，表示训练过程中不使用 Dice Loss 作为损失函数，而可能会使用其他损失函数，比如交叉熵损失函数或者 Focal Loss 等。当 dice_loss 设置为 True 时，表示使用 Dice Loss 作为损失函数进行训练。通常，对于语义分割任务，Dice Loss 是一种常用的损失函数，可以有效地提高模型的分割性能。
建议设置：
建议根据种类数和批量大小来选择是否使用 Dice Loss 的原因涉及到损失函数在训练过程中的表现和计算复杂度。
-
种类少（几类）时，设置为 True：
- 当类别较少时，每个像素的分类任务相对简单，使用 Dice Loss 可以很好地对模型进行优化，因为 Dice Loss 在处理二分类问题时效果较好。这样可以帮助模型更好地区分目标和背景，有助于提高模型的分割精度。
-
种类多（十几类）时，如果 batch_size 比较大（10 以上），那么设置为 True：
- 当类别较多且批量大小较大时，每个批次中包含了大量的样本和类别信息。Dice Loss 在这种情况下也能表现得比较好，因为在大批量数据中，Dice Loss 可以更好地优化网络参数，提高模型的收敛速度和效果。
-
种类多（十几类）时，如果 batch_size 比较小（10 以下），那么设置为 False：
- 当类别较多且批量大小较小时，每个批次中的样本和类别信息相对有限。Dice Loss 可能会受到噪声影响，导致模型训练不稳定，甚至可能会陷入局部最优点。在这种情况下，使用其他损失函数，如交叉熵损失函数或 Focal Loss，可能更适合，因为它们在处理多分类问题时更具优势。
总体来说，Dice Loss 适用于二分类任务，在分割问题中特别适用于轮廓明显、目标与背景差异较大的情况。对于多分类问题，尤其是当类别较多且批量大小较小时，可以考虑使用其他损失函数，根据实际情况选择最适合的损失函数有助于提高模型的训练效果。
focal_loss = False # 是否使用focal loss来防止正负样本不平衡
cls_weights = np.ones([num_classes], np.float32)
focal_los
是否给不同种类赋予不同的损失权值，默认是平衡的。
设置的话，注意设置成numpy形式的，长度和num_classes一样。
如：
num_classes = 3
cls_weights = np.array([1, 2, 3], np.float32)
num_workers = 4
num_workers
用于设置是否使用多线程读取数据，1代表关闭多线程
开启后会加快数据读取速度，但是会占用更多内存
keras里开启多线程有些时候速度反而慢了许多
在IO为瓶颈的时候再开启多线程，即GPU运算速度远大于读取图片的速度。
num_workers参数用于设置数据加载过程中使用的多线程数量。在 PyTorch 的DataLoader中，可以通过设置num_workers参数来启用多线程数据加载，从而加快数据读取速度。每个工作线程负责预加载数据，使得 CPU 在读取数据的同时可以进行其他操作，从而提高数据加载的效率。
设置 num_workers 大于1时，会开启多个工作线程来加载数据，这样可以在一定程度上加快数据加载速度，特别是在数据读取过程中是瓶颈的情况下，如硬盘 IO 较慢或数据集较大的情况下。
但是，开启多线程也会带来一些额外的开销和内存消耗。如果 num_workers 设置得太大，可能会导致内存占用过多，进而影响模型训练的稳定性。此外，多线程的效果也取决于硬件环境和数据集的大小，有时候可能会因为线程间的竞争或数据集本身较小而导致速度没有明显的提升，甚至可能导致速度下降。
在具体使用时，可以根据硬件配置和数据集大小进行调优。如果 GPU 的运算速度远大于数据读取速度，即数据读取是瓶颈，开启多线程可以提高数据加载效率。但如果数据读取速度和 GPU 运算速度相当，或者数据集较小，开启多线程可能不会带来明显的速度提升，甚至可能因为多线程带来的开销而导致速度下降。因此，最佳的 num_workers 设置应该在实际使用中根据情况进行调整。
训练准备
设置用到的显卡
ngpus_per_node = torch.cuda.device_count() # 当前节点上可用的GPU数量
if distributed: # 检查代码是否在分布式训练环境中运行。
dist.init_process_group(backend="nccl")
"""
如果代码在分布式环境中运行，这行代码会初始化进程组，以便不同进程之间进行通信。
使用nccl后端用于在单个节点内的GPU之间进行通信。
"""
local_rank = int(os.environ["LOCAL_RANK"])
"""
它从环境变量中提取当前进程的本地排名（local rank）。
在分布式训练中，每个进程被分配一个本地排名，对应于它将用于计算的GPU。
"""
rank = int(os.environ["RANK"])
"""从环境变量中提取当前进程的全局排名（global rank）。
全局排名是分布式设置中每个进程的唯一标识符。
"""
device = torch.device("cuda", local_rank)
"""
根据本地排名创建一个PyTorch设备对象。
如果代码在分布式环境中运行，则设备将根据本地排名设置为相应的GPU。
否则，它会检查是否有支持CUDA的GPU，并将设备设置为'cuda'，否则设置为'cpu'。
"""
if local_rank == 0:
"""
用于只从本地排名为0的进程（通常是第一个GPU）打印一些信息。
它显示了rank和local_rank，表示哪个进程在哪个GPU上运行，以及可用的总GPU数量。
"""
print(f"[{os.getpid()}] (rank = {rank}, local_rank = {local_rank}) training...") #
print("Gpu Device Count : ", ngpus_per_node) #
else:
"""
如果代码不在分布式环境中运行，则执行此代码块。
在这种情况下，它检查GPU是否可用，如果可用，则将设备设置为'cuda'，否则设置为'cpu'。
此外，在这种情况下，local_rank被设置为0，因为只有一个进程（非分布式）。
"""
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') #
local_rank = 0 #
下载预训练权重
#首先，检查是否需要加载预训练权重。如果pretrained=True，继续进行权重加载和初始化；否则跳过这部分操作。
if pretrained:
# 如果在分布式训练中（distributed=True），根据本地进程的local_rank判断是否为rank为0的进程。
# 只有rank为0的进程才会执行权重下载操作download_weights(phi)，其他进程等待该操作完成。
if distributed:
if local_rank == 0:
download_weights(phi)
dist.barrier()
#如果不是分布式训练或者是rank为0的进程，也会执行权重下载操作。
else:
download_weights(phi)
# 创建一个模型SegFormer，并传入一些参数：num_classes（类别数量）、phi（用于选择模型的大小）、pretrained（是否使用预训练权重）。
model = SegFormer(num_classes=num_classes, phi=phi, pretrained=pretrained)
# 如果pretrained为False，调用weights_init(model)对模型进行权重的初始化。
if not pretrained:
weights_init(model)
# 如果指定了model_path（模型权重文件的路径），开始加载预训练权重。
if model_path != '':
# 如果是rank为0的进程，显示加载权重的提示信息。
if local_rank == 0:
print('Load weights {}.'.format(model_path))
#------------------------------------------------------#
# 根据预训练权重的Key和模型的Key进行加载
#------------------------------------------------------#
# 获取当前模型的权重字典model_dict。
model_dict = model.state_dict()
# 使用torch.load加载预训练模型的权重字典pretrained_dict，并将其加载到指定设备上（GPU或CPU）。
pretrained_dict = torch.load(model_path, map_location = device)
# 定义三个列表，用于记录加载成功匹配的键、未成功匹配的键以及临时字典。
load_key, no_load_key, temp_dict = [], [], {}
# 遍历预训练模型的权重字典，与当前模型的权重字典进行对比。
for k, v in pretrained_dict.items():
# 如果键k在当前模型的权重字典中，并且形状也匹配，将权重添加到临时字典temp_dict中，并记录到load_key中。
if k in model_dict.keys() and np.shape(model_dict[k]) == np.shape(v):
temp_dict[k] = v
load_key.append(k)
else:
no_load_key.append(k)
model_dict.update(temp_dict)
model.load_state_dict(model_dict)
#------------------------------------------------------#
# 显示没有匹配上的Key
#------------------------------------------------------#
# 显示加载结果：成功匹配和未匹配的权重键，并提示用户可能出现的问题。
if local_rank == 0:
print("\nSuccessful Load Key:", str(load_key)[:500], "……\nSuccessful Load Key Num:", len(load_key))
print("\nFail To Load Key:", str(no_load_key)[:500], "……\nFail To Load Key num:", len(no_load_key))
print("\n\033[1;33;44m温馨提示，head部分没有载入是正常现象，Backbone部分没有载入是错误的。\033[0m")
"根据预训练权重的Key和模型的Key进行加载"的意思是将预训练模型的权重加载到当前模型中，但要根据它们的键（Key）来匹配。
在深度学习模型中，每个层或参数都有一个唯一的键（Key），用于在字典或类似的数据结构中标识它们。这些键通常是字符串，用于存储和获取模型的权重。
当我们从预训练模型加载权重时，预训练模型和当前模型的架构可能有所不同，因此它们的权重的键可能不完全匹配。为了正确加载权重，我们需要通过对比键来找到匹配的权重。
具体地，这段代码中的操作是将预训练模型的权重字典pretrained_dict加载到当前模型的权重字典model_dict中。通过循环遍历预训练模型的权重字典，对于每个键和值，检查它是否存在于当前模型的权重字典中，并且与当前模型的权重形状匹配（大小一致）。如果匹配成功，则将该权重添加到当前模型的权重字典中。
这样做的目的是确保只有那些在当前模型中存在且形状匹配的权重被加载进来，从而避免由于权重形状不匹配而导致的错误。没有匹配的键将被忽略，这可能是因为当前模型结构中不存在该层或参数。
这个过程通常在迁移学习中使用，以便从预训练模型中初始化当前模型的权重，从而加速模型的训练和提高模型性能。
训练
记录Loss
if local_rank == 0:
time_str = datetime.datetime.strftime(datetime.datetime.now(),'%Y_%m_%d_%H_%M_%S')
log_dir = os.path.join(save_dir, "loss_" + str(time_str))
loss_history = LossHistory(log_dir, model, input_shape=input_shape)
else:
loss_history = None
if fp16:
from torch.cuda.amp import GradScaler as GradScaler
scaler = GradScaler()
else:
scaler = None
model_train = model.train()
torch 1.2不支持amp，建议使用torch 1.7.1及以上正确使用fp16
因此torch1.2这里显示"could not be resolve"
多卡同步Bn
if sync_bn and ngpus_per_node > 1 and distributed:
model_train = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model_train)
elif sync_bn:
print("Sync_bn is not support in one gpu or not distributed.")
if Cuda:
if distributed:
#----------------------------#
# 多卡平行运行
#----------------------------#
model_train = model_train.cuda(local_rank)
model_train = torch.nn.parallel.DistributedDataParallel(model_train, device_ids=[local_rank],
find_unused_parameters=True)
else:
model_train = torch.nn.DataParallel(model)
cudnn.benchmark = True
model_train = model_train.cuda()
多卡同步批归一化（Sync Batch Normalization）是在多个GPU上进行训练时的一种优化技术，主要用于加速训练并提高模型的收敛性。
在深度学习中，批归一化（Batch Normalization，简称BN）用于加速训练过程，稳定梯度更新，从而加快模型的收敛速度。然而，当使用多个GPU并行训练时，由于每个GPU的批次大小可能不同，传统的批归一化会出现问题。多卡同步批归一化通过在多个GPU之间同步均值和方差信息，使得在每个GPU上计算的均值和方差都是整个训练批次的统计值，从而保证了每个GPU上的批归一化操作的一致性。
读取数据集对应的txt
with open(os.path.join(VOCdevkit_path, "VOC2007/ImageSets/Segmentation/train.txt"),"r") as f:
train_lines = f.readlines()
with open(os.path.join(VOCdevkit_path, "VOC2007/ImageSets/Segmentation/val.txt"),"r") as f:
val_lines = f.readlines()
num_train = len(train_lines)
num_val = len(val_lines)
)
显示训练配置信息
if local_rank == 0:
# 显示训练配置
show_config(
num_classes = num_classes, phi = phi, model_path = model_path,
input_shape = input_shape, \
Init_Epoch = Init_Epoch, Freeze_Epoch = Freeze_Epoch,
UnFreeze_Epoch = UnFreeze_Epoch, Freeze_batch_size = Freeze_batch_size,
Unfreeze_batch_size = Unfreeze_batch_size, Freeze_Train = Freeze_Train, \
Init_lr = Init_lr, Min_lr = Min_lr, optimizer_type = optimizer_type,
momentum = momentum, lr_decay_type = lr_decay_type, \
save_period = save_period, save_dir = save_dir, num_workers = num_workers,
num_train = num_train, num_val = num_val
)
配置epoch和patch size
"""
`wanted_step`和`total_step`：分别表示预期的总训练步长和实际计算得到的总训练步长。
其中，`wanted_step`是根据优化器类型（optimizer_type）来确定的，`total_step`通过计算得到，
其值为解冻阶段的总训练步长。
"""
wanted_step = 1.5e4 if optimizer_type == "adamw" else 0.5e4
total_step = num_train // Unfreeze_batch_size * UnFreeze_Epoch
#如果实际的总训练步长小于等于预期的总训练步长。
if total_step <= wanted_step:
#如果解冻阶段的训练集样本数除以解冻阶段的批量大小为0（即训练集样本数太少）。
if num_train // Unfreeze_batch_size == 0:
raise ValueError('数据集过小，无法进行训练，请扩充数据集。')
# 根据预期的总训练步长和解冻阶段的批量大小计算出建议的总训练世代数。
wanted_epoch = wanted_step // (num_train // Unfreeze_batch_size) + 1
#打印建议信息：根据计算结果打印建议的训练步长和总训练世代数。这些信息对于训练的调优和合理设置非常有帮助。
print("\n\033[1;33;44m[Warning] 使用%s优化器时，建议将训练总步长设置到%d以上。\033[0m"%(optimizer_type, wanted_step))
print("\033[1;33;44m[Warning] 本次运行的总训练数据量为%d，Unfreeze_batch_size为%d，共训练%d个Epoch，计算出总训练步长为%d。\033[0m"%(num_train, Unfreeze_batch_size, UnFreeze_Epoch, total_step))
print("\033[1;33;44m[Warning] 由于总训练步长为%d，小于建议总步长%d，建议设置总世代为%d。\033[0m"%(total_step, wanted_step, wanted_epoch))
总训练世代指的是遍历全部数据的总次数 总训练步长指的是梯度下降的总次数
每个训练世代包含若干训练步长，每个训练步长进行一次梯度下降。
此处仅建议最低训练世代，上不封顶，计算时只考虑了解冻部分
冻结or不冻结
if True:
UnFreeze_flag = False
#------------------------------------#
# 冻结一定部分训练
#------------------------------------#
if Freeze_Train:
for param in model.backbone.parameters():#遍历模型的backbone（骨干网络）参数。
# 将backbone参数的requires_grad属性设置为False，表示这些参数在训练过程中不需要梯度更新，
# 即进行参数冻结。
param.requires_grad = False
#-------------------------------------------------------------------#
# 如果不冻结训练的话，直接设置batch_size为Unfreeze_batch_size
#-------------------------------------------------------------------#
batch_size = Freeze_batch_size if Freeze_Train else Unfreeze_batch_size
#-------------------------------------------------------------------#
# 判断当前batch_size，自适应调整学习率
#-------------------------------------------------------------------#
nbs = 16
lr_limit_max = 1e-4 if optimizer_type in ['adam', 'adamw'] else 5e-2
lr_limit_min = 3e-5 if optimizer_type in ['adam', 'adamw'] else 5e-4
Init_lr_fit = min(max(batch_size / nbs * Init_lr, lr_limit_min), lr_limit_max)
Min_lr_fit = min(max(batch_size / nbs * Min_lr, lr_limit_min * 1e-2), lr_limit_max * 1e-2)
#---------------------------------------#
# 根据optimizer_type选择优化器
#---------------------------------------#
optimizer = {
'adam' : optim.Adam(model.parameters(), Init_lr_fit, betas = (momentum, 0.999),
weight_decay = weight_decay),
'adamw' : optim.AdamW(model.parameters(), Init_lr_fit, betas = (momentum, 0.999),
weight_decay = weight_decay),
'sgd' : optim.SGD(model.parameters(), Init_lr_fit, momentum = momentum, nesterov=True,
weight_decay = weight_decay)
}[optimizer_type]
#---------------------------------------#
# 获得学习率下降的公式
#---------------------------------------#
lr_scheduler_func = get_lr_scheduler(lr_decay_type, Init_lr_fit, Min_lr_fit, UnFreeze_Epoch)
#---------------------------------------#
# 判断每一个世代的长度
#---------------------------------------#
epoch_step = num_train // batch_size
epoch_step_val = num_val // batch_size
if epoch_step == 0 or epoch_step_val == 0:
raise ValueError("数据集过小，无法继续进行训练，请扩充数据集。")
train_dataset = SegmentationDataset(train_lines, input_shape, num_classes, True, VOCdevkit_path)
val_dataset = SegmentationDataset(val_lines, input_shape, num_classes, False, VOCdevkit_path)
if distributed:
train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset, shuffle=True,)
val_sampler = torch.utils.data.distributed.DistributedSampler(val_dataset, shuffle=False,)
batch_size = batch_size // ngpus_per_node
shuffle = False
else:
train_sampler = None
val_sampler = None
shuffle = True
gen = DataLoader(train_dataset, shuffle = shuffle, batch_size = batch_size, num_workers = num_workers,
pin_memory=True,
drop_last = True, collate_fn = seg_dataset_collate, sampler=train_sampler)
gen_val = DataLoader(val_dataset , shuffle = shuffle, batch_size = batch_size,
num_workers = num_workers, pin_memory=True,
drop_last = True, collate_fn = seg_dataset_collate, sampler=val_sampler)
#----------------------#
# 记录eval的map曲线
#----------------------#
if local_rank == 0:
eval_callback = EvalCallback(model, input_shape, num_classes, val_lines, VOCdevkit_path, log_dir, Cuda, \
eval_flag=eval_flag, period=eval_period)
else:
eval_callback = None
开始模型训练
for epoch in range(Init_Epoch, UnFreeze_Epoch):
#---------------------------------------#
# 如果模型有冻结学习部分
# 则解冻，并设置参数
#---------------------------------------#
if epoch >= Freeze_Epoch and not UnFreeze_flag and Freeze_Train:
batch_size = Unfreeze_batch_size
#-------------------------------------------------------------------#
# 判断当前batch_size，自适应调整学习率
#-------------------------------------------------------------------#
nbs = 16
lr_limit_max = 1e-4 if optimizer_type in ['adam', 'adamw'] else 5e-2
lr_limit_min = 3e-5 if optimizer_type in ['adam', 'adamw'] else 5e-4
Init_lr_fit = min(max(batch_size / nbs * Init_lr, lr_limit_min), lr_limit_max)
Min_lr_fit = min(max(batch_size / nbs * Min_lr, lr_limit_min * 1e-2), lr_limit_max * 1e-2)
#---------------------------------------#
# 获得学习率下降的公式
#---------------------------------------#
lr_scheduler_func = get_lr_scheduler(lr_decay_type, Init_lr_fit, Min_lr_fit, UnFreeze_Epoch)
for param in model.backbone.parameters():
param.requires_grad = True
epoch_step = num_train // batch_size
epoch_step_val = num_val // batch_size
if epoch_step == 0 or epoch_step_val == 0:
raise ValueError("数据集过小，无法继续进行训练，请扩充数据集。")
gen = DataLoader(train_dataset, shuffle = shuffle, batch_size = batch_size,
num_workers = num_workers, pin_memory=True,
drop_last = True, collate_fn = seg_dataset_collate,
sampler=train_sampler)
gen_val = DataLoader(val_dataset , shuffle = shuffle, batch_size = batch_size,
num_workers = num_workers, pin_memory=True,
drop_last = True, collate_fn = seg_dataset_collate, sampler=val_sampler)
UnFreeze_flag = True
if distributed:
train_sampler.set_epoch(epoch)
set_optimizer_lr(optimizer, lr_scheduler_func, epoch)
fit_one_epoch(model_train, model, loss_history, eval_callback, optimizer, epoch, epoch_step, epoch_step_val,
gen, gen_val, UnFreeze_Epoch, Cuda, \
dice_loss, focal_loss, cls_weights, num_classes, fp16, scaler, save_period, save_dir, local_rank)
if distributed:
dist.barrier()
if local_rank == 0:
loss_history.writer.close()
扫一扫
3万+
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
