项目代码: 来自github的YOLOv1开源项目
本文是关于train.py的详细理解。
-
该文件一共有三个函数
- parse_args()
- set_lr(optimizer, lr)
- train()
parse_args()：
def parse_args():
parser = argparse.ArgumentParser(description='YOLO Detection')
parser.add_argument('-v', '--version', default='yolo',
help='yolo')
parser.add_argument('-hr', '--high_resolution', action='store_true', default=False,
help='use high resolution to pretrain.')
return parser.parse_args()
- 其中第一个函数创建了argparse模块，这个模块用于解析命令行的输入。原本修改参数需要通过修改python文件，通过使用这个模块，就可以在命令行修改参数，方便了很多。
- ArgumentParser是argparse模块中的一个类，它提供了一种解析命令行参数的方式。这里将这个类实例化为parser，并且在实例化这个类的时候写了一个description参数，在cmd窗口使用-h选项的时候能够显示这里的description参数的值，用于提示用户当前程序是做什么应用的。
- add_argument是ArgumentParser这个类的一个方法，目的是添加参数，一个-代表的是短选项名字，- -代表的是长选项名字，v这个参数的意思是选择版本，该代码中只有yolo可用，作者在train方法中写，当该参数不为yolo的时候会
print('We only support YOLO !!!')。 - action参数设置为store_true表示，如果在命令行中写了该命令，那么该参数的值就会被设置成true，输入-hr之后，args.hr的值就为true了，在接下来的程序中，就能够通过
if args.hr来判断用户是否需要使用高分辨率。 return parser.parse_args()这里parser实例调用了parse_args()方法，该方法能够解析命令行的输入，返回一个包含命令行参数值的命名对象空间。train方法中写了args = parse_args()，这里args的值就是很多个参数和其对应的值的集合，可以通过args.参数名来调用。
set_lr(optimizer, lr):
def set_lr(optimizer, lr):
for param_group in optimizer.param_groups:
param_group['lr'] = lr
- 这个函数的目的是，统一设置学习率，作者在这里进行了一下封装。optimizer是优化器，lr是学习率。
optimizer.param_groups是一个包含所有参数组的列表，其中每个参数组都是一个字典，每个字典中包含很多键值对。- 举个例子
# 创建一个简单的线性模型
model = nn.Linear(10, 1)
# 将权重和偏置分别划分成两个参数组
params = [
{"params": model.weight, "lr": 0.01},
{"params": model.bias, "lr": 0.1},
]
# 创建优化器，设置默认学习率为0.01，但是bias的学习率还是0.1
optimizer = SGD(params, lr=0.01)
train()：
def train():
args = parse_args()
- 这里是读取命令行中用户的传参。
path_to_save = os.path.join(
args.save_folder,
args.dataset,
args.version)
os.makedirs(path_to_save, exist_ok=True)
- os.path.join方法是将这几个参数按照路径的方式连接起来，然后用os.makedirs方法创建这个目录，
exist_ok=True有了这个参数后，如果该目录已经存在，也不会报错。
# use hi-res backbone
if args.high_resolution:
print('use hi-res backbone')
hr = True
else:
hr = False
- 如果命令行中输入了
-hr，根据上面的代码，这里的args.high_resolution的值就是true。
# cuda
if args.cuda:
print('use cuda')
cudnn.benchmark = True
device = torch.device("cuda")
else:
device = torch.device("cpu")
cudnn.benchmark = True是将 PyTorch 中的 cuDNN 加速库设置为使用自动寻找最优配置的模式，从而加速卷积操作。cuDNN 加速库会在第一次执行卷积操作时，自动寻找最优配置，并将结果缓存下来，从而在后续卷积操作中提高性能，这样的话第一次卷积操作的性能也会受到影响。
# multi-scale
if args.multi_scale:
print('use the multi-scale trick ...')
train_size = [640, 640]
val_size = [416, 416]
else:
train_size = [416, 416]
val_size = [416, 416]
- multi-scale 训练的基本思想是在训练过程中随机使用多个不同尺度的输入图像，从而增加模型对尺度变化的鲁棒性，提高模型的泛化能力。这里如果采用multi-scale训练，训练集和验证集的图片尺寸就会设置成不一样的。
cfg = train_cfg
- 这里的
train_cfg是data文件下的config.py里的字典具体内容如下：
train_cfg = {
'lr_epoch': (60, 90, 160),
'max_epoch': 160,
'min_dim': [416, 416]
}
- lr_epoch: 是一个元组，表示在第 60、90、160 个 epoch 处要调整学习率，max_epoch: 表示最大的训练 epoch 数量，即模型要训练多少个 epoch，min_dim: 是一个列表，表示训练时输入图像的最小尺寸。
if args.dataset == 'voc':
data_dir = VOC_ROOT
num_classes = 20
dataset = VOCDetection(root=data_dir,
img_size=train_size[0],
transform=SSDAugmentation(train_size)
)
evaluator = VOCAPIEvaluator(data_root=data_dir,
img_size=val_size,
device=device,
transform=BaseTransform(val_size),
labelmap=VOC_CLASSES
)
- VOCDetection继承了data.Dataset类，用于加载训练数据集。
VOC_ROOT是voc0712.py文件中的一个路径，这里也可以写VOCdevkit文件夹所在的路径。这里的img_size是图片的大小，transform是指定如何增强数据集的一个参数，这里使用的是空间金字塔池化SSD (Single Shot MultiBox Detector）来处理输入图像的不同尺度信息，以提高目标检测的准确性。具体来说，SSD网络通过将输入图像分别经过不同大小的卷积核进行卷积，然后通过空间金字塔池化将不同尺度的特征图合并起来，从而获得全局的特征信息。 evaluator是用于评估模型性能的对象，它是一个VOCAPIEvaluator类的实例。evaluator对象可以通过evaluate方法，评估传入参数 model 表示的模型的性能，即evaluator.evaluate(model)。VOC_CLASSES是一个包含20个字符串的Python列表，用于表示VOC数据集中的20个不同的对象类别。在模型中，通常会将不同类别映射到数字编码，比如将"cat"映射到数字0，VOC_CLASSES列表将类别名称与数字编码对应起来，使得可以将类别名称转换为数字编码，以便于在模型中进行处理。labelmap参数规定类别名称和数字的一个映射。device的值是device，在上文cuda参数的处理那里设置好了。
dataloader = torch.utils.data.DataLoader(
`dataset`,
batch_size=args.batch_size,
shuffle=True,
collate_fn=detection_collate,
num_workers=args.num_workers,
pin_memory=True
)
- DataLoader这个类的作用是读取数据，这里将其实例化为dataloader，参数
dataset是上文的数据集，batch_size是每次训练的样本数量，collate_fn是将样本列表转换为 mini-batch 张量的一个可选参数，detection_collate对于每个图像，将其变换为相同的大小，同时将它们堆叠在一起形成一个 4D 张量（即，大小为 (batch_size, channel, height, width) 的张量），并将其返回。它还将每个样本中的标注框与其对应的图像一起组合在一个列表中返回。num_workers参数用于指定使用多少个子进程来加载数据。
if args.version == 'yolo':
from models.yolo import myYOLO
yolo_net = myYOLO(device, input_size=train_size, num_classes=num_classes, trainable=True)
print('Let us train yolo on the %s dataset ......' % (args.dataset))
model = yolo_net
model.to(device).train()
- 实例化作者写的YOLO模型，
model.to(device)的作用是将模型的参数和缓冲区移动到指定的设备上，而model.train()的作用是将模型的training属性设置为True，这将启用模型中的训练特定的操作，例如Dropout和BatchNorm等。使用GPU进行训练时，需要将模型、数据和计算所需的其他变量都放到GPU上。
# use tfboard
if args.tfboard:
print('use tensorboard')
from torch.utils.tensorboard import SummaryWriter
c_time = time.strftime('%Y-%m-%d %H:%M:%S',time.localtime(time.time()))
log_path = os.path.join('log/coco/', args.version, c_time)
os.makedirs(log_path, exist_ok=True)
writer = SummaryWriter(log_path)
c_time这行代码是获取当前时间，转换为当地时间，并将其格式化为指定格式的字符串。SummaryWriter是TensorBoard的PyTorch接口。
if args.resume is not None:
print('keep training model: %s' % (args.resume))
model.load_state_dict(torch.load(args.resume, map_location=device))
- 这里代码的意思是，是不是接着上次训练的过程继续训练。
load_state_dict方法用于加载一个预训练模型的参数到当前模型中。torch.load函数会读取已保存的模型参数，返回值是一个包含加载对象的字典或者其他数据结构，args.resume是预训练模型的路径，map_location=device表示将预训练模型加载到当前的设备（CPU或GPU）中。
# optimizer setup
base_lr = args.lr
tmp_lr = base_lr
optimizer = optim.SGD(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay
)
max_epoch = cfg['max_epoch']
epoch_size = len(dataset) // args.batch_size
optimizer是一个实例化后的优化器，优化器的作用是在每个batch结束后就用梯度下降法更新参数的值，model.parameters()返回一个包含模型所有参数的迭代器，可以用于更新模型的参数。具体来说，它返回一个包含模型权重和偏置等参数的列表，这些参数是需要在训练过程中进行优化的。lr是学习率，控制参数更新的幅度，momentum是一个超参数，用于控制梯度下降算法在更新参数时的加速度，每次更新时，梯度下降算法会将当前梯度的一个分量（即当前步骤的方向）与上一次的分量进行加权平均，从而加速收敛速度。weight_decay参数是正则化系数，它控制正则化惩罚的强度，对应于正则化项中的λ值。max_epoch是使用者设置的参数，epoch_size是根据batchsize算出来的数。
for epoch in range(args.start_epoch, max_epoch):
- 这里
start_epoch可能是接着之前的训练epoch继续训练。
# use cos lr
if args.cos and epoch > 20 and epoch <= max_epoch - 20:
# use cos lr
tmp_lr = 0.00001 + 0.5*(base_lr-0.00001)*(1+math.cos(math.pi*(epoch-20)*1./ (max_epoch-20)))
set_lr(optimizer, tmp_lr)
elif args.cos and epoch > max_epoch - 20:
tmp_lr = 0.00001
set_lr(optimizer, tmp_lr)
# use step lr
else:
if epoch in cfg['lr_epoch']:
tmp_lr = tmp_lr * 0.1
set_lr(optimizer, tmp_lr)
- 这里是说，如果选了coslr，也就是学习率余弦退火，余弦退火法能够在训练初期较快地下降学习率，然后在训练后期逐渐降低学习率，以提高模型收敛速度和稳定性。最后20个epoch保持学习率为0.00001，在距离训练结束还有一定的epoch数时，将学习率降低到一个较小的值可以帮助模型更加充分地探索局部最优解空间，并减小因学习率过大而错过最优解的风险。
for iter_i, (images, targets) in enumerate(dataloader):
# WarmUp strategy for learning rate
if not args.no_warm_up:
if epoch < args.wp_epoch:
tmp_lr = base_lr * pow((iter_i+epoch*epoch_size)*1. / (args.wp_epoch*epoch_size), 4)
# tmp_lr = 1e-6 + (base_lr-1e-6) * (iter_i+epoch*epoch_size) / (epoch_size * (args.wp_epoch))
set_lr(optimizer, tmp_lr)
elif epoch == args.wp_epoch and iter_i == 0:
tmp_lr = base_lr
set_lr(optimizer, tmp_lr)
enumerate()常用于在循环中同时获得每个元素的值和该元素的索引，返回一个包含每个元素的索引和该元素的元组，可以用在 for 循环中对序列进行迭代。WarmUp策略指的是在模型训练的前期使用比较小的学习率进行预热，使得模型更好地适应新的任务，减少训练的波动性，从而提高训练的效果。((iter_i+epoch*epoch_size)*1. / (args.wp_epoch*epoch_size))，这个是warmup系数，该数字是从0到1增加的，乘以1.0是为了将整数转换为浮点数，以便在除法运算时获得精确的小数结果。epoch_size是每个 epoch 中的迭代次数，即每个 epoch 中有多少个 mini-batch，args.wp_epoch是 WarmUp 阶段的 epoch 数。iter_i+epoch*epoch_size表示当前batch在整个训练数据中的位置。args.wp_epoch*epoch_size是warmup阶段总共需要训练的batch数目。
# multi-scale trick
if iter_i % 10 == 0 and iter_i > 0 and args.multi_scale:
# randomly choose a new size
size = random.randint(10, 19) * 32
train_size = [size, size]
model.set_grid(train_size)
if args.multi_scale:
# interpolate
images = torch.nn.functional.interpolate(images, size=train_size, mode='bilinear', align_corners=False)
- 这里是多尺度训练技巧，目的是为了提高模型检测不同大小的物体的能力。当迭代次数是10的倍数时，就随机选择一个新的大小，增加训练样本的多样性。
random.randint(10, 19)返回一个10-19之间的随机整数，然后再乘以32，得到320-608之间的随机数，用作新的图片大小。model.set_grid(train_size)是yolo文件的一个方法，根据输入的图片尺寸创建网格，具体见yolo.py代码理解。torch.nn.functional.interpolate是一个对张量进行插值操作的函数，这里images通过双线性插值的方式调整到指定的train_size大小。
# make train label
targets = [label.tolist() for label in targets]
targets = tools.gt_creator(input_size=train_size, stride=yolo_net.stride, label_lists=targets)
targets = torch.tensor(targets).float().to(device)
- 这段代码是将原始的标签转换成模型训练需要的格式，即生成真实的目标检测的ground truth label，
tools.gt_creator返回值为Numpy数组，然后转换为tensor类型。具体见tools.py文件的理解。
conf_loss, cls_loss, txtytwth_loss, total_loss = model(images, target=targets)
- 在调用
model(images, target=targets)时，实际上是调用了model.__call__(images, target=targets)，这个方法是定义好的，在这里__call__方法会调用forward方法，model()和model.__call__()是一样的。
# backprop
total_loss.backward()
optimizer.step()
optimizer.zero_grad()
- 这里是反向传播，套路化步骤，包括计算梯度和反向传播，以及更新优化器参数和清零梯度。
total_loss.backward()通过调用backward()方法自动计算每个可训练参数的梯度，并将其存储在相应的.grad属性中。optimizer.step()根据计算出的梯度更新每个可训练参数的值，即执行优化器的更新操作。optimizer.zero_grad()清空可训练参数的梯度，以便进行下一轮迭代的计算，这里的梯度是用一个batch算出来的梯度。
# display
if iter_i % 10 == 0:
if args.tfboard:
# viz loss
writer.add_scalar('object loss', conf_loss.item(), iter_i + epoch * epoch_size)
writer.add_scalar('class loss', cls_loss.item(), iter_i + epoch * epoch_size)
writer.add_scalar('local loss', txtytwth_loss.item(), iter_i + epoch * epoch_size)
t1 = time.time()
print('[Epoch %d/%d][Iter %d/%d][lr %.6f]'
'[Loss: obj %.2f || cls %.2f || bbox %.2f || total %.2f || size %d || time: %.2f]'
% (epoch+1, max_epoch, iter_i, epoch_size, tmp_lr,
conf_loss.item(), cls_loss.item(), txtytwth_loss.item(), total_loss.item(), train_size[0], t1-t0),
flush=True)
- 这段代码是每训练10个iteration打印一次当前的loss和运行时间，并将loss写入tensorboard的可视化界面中。
writer.add_scalar是用于将标量数据写入tensorboard的函数。该函数需要三个参数：tag、value和step。其中，tag是指标签，通常表示要绘制的曲线的名称；value是指要记录的标量数据；step表示训练的步数或迭代次数，用于在x轴上显示。writer.add_scalar需要的是一个标量值，因此需要通过.item()方法将Tensor对象转换为 Python 数值。iter_i + epoch * epoch_size计算出当前的迭代次数。后面print出了一些训练过程的结果。
# evaluation
if (epoch + 1) % args.eval_epoch == 0:
model.trainable = False
model.set_grid(val_size)
model.eval()#在评估模式下，模型的行为与训练模式下不同，它会禁用一些特定的操作，如 dropout 和 batch normalization。
# evaluate
evaluator.evaluate(model)
# convert to training mode.
model.trainable = True
model.set_grid(train_size)
model.train()
args.eval_epoch表示每训练多少个epoch后进行一次模型评估，这里+1是因为epoch是从0开始算的。将模型的trainable属性设置为False，表示不参与反向传播，不更新权重参数，这样做是为了保证在评估时模型的权重不会发生改变，结果也更加准确。通过调用model.set_grid(val_size)方法将输入大小设置为验证集大小。将模型的运行模式设置为评估模式，即model.eval()，这样可以禁用一些特定的操作，如dropout和batch normalization。调用评估器（evaluator）的evaluate方法，对模型进行评估，评估函数将计算一些性能指标（如准确率、精确率、召回率等）并将它们输出到控制台或记录到文件中。将模型的trainable属性重新设置为True，将输入大小设置为训练集大小，将模型的运行模式设置为训练模式，即model.train()。
# save model
if (epoch + 1) % 10 == 0:
print('Saving state, epoch:', epoch + 1)
torch.save(model.state_dict(), os.path.join(path_to_save,
args.version + '_' + repr(epoch + 1) + '.pth')
)
- 这段代码用于在每10个epoch结束时保存模型的权重。
model.state_dict()返回模型的所有权重参数，并使用torch.save函数将其保存到指定路径下的pth文件中，args.version + '_' + repr(epoch + 1) + '.pth'则是文件的命名方式，repr(epoch+1)将整型epoch+1转换为字符串类型，在这里repr函数和str函数没有区别。 但是，在其他情况下，repr 和 str 之间有一些区别。 repr 返回一个对象的完整描述，通常可以通过 eval() 函数恢复这个对象，而 str 返回一个更简单的字符串表示，用于更常规的显示目的。
if __name__ == '__main__':
train()
- 最后在主函数中调用train函数，即这个文件的入口。
470
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
