源码文件
./maskrcnn_benchmark/engine/trainer.py
./maskrcnn_benchmark/engine/inference.py
模型测试/推演时, 无需计算每个参数的梯度, 因此会大大减少 GPU 的显存占用空间
trainer.py 文件概览
本文件中的代码是用来定义模型训练的步骤的, 当我们确定了模型, 数据集载入器, 优化器, 更新策略等相关组件以后, 就可以调用该文件中的函数来进行模型训练, 下面先简单看一下本文件的函数概览.
1 | # ./maskrcnn_benchmark/engine/trainer.py |
trainer 导入各种包及函数
下面是本文件导入的包及函数, 我们会对这些包进行简要的介绍, 并且会给出详细解析的博文链接.
1 | # ./maskrcnn_benchmark/engine/trainer.py |
get_world_size
: 该函数封装了 troch.distributed.get_world_size()
函数, 位于 ./maskrcnn_benchmark/utils/comm.py
文件中, 详细解析请看commMetricLogger
: 该类中包含了同文件的 SmoothedValue
类作为其数据成员, SmoothedValue
类用于跟踪一系列的值, 同时还会提供访问这些值的滑动平均值或全局平均值, 而位于 MetricLogger
定义了将这些值打印到屏幕或保存到文件中的代码, 这两个类位于 ./maskrcnn_benchmark/utils.metric_logger
文件中, 详细解析请看metric_logger
trainer.do_train() 模型训练核心函数
我们按照函数的一般调用顺序, 首先来看看模型训练的核心逻辑代码的实现, 具体如下:
1 | # ./maskrcnn_benchmark/engine/trainer.py |
trainer.reduce_loss_dict() 计算 reduced_loss
下面的代码会计算多 GPU 时的 reduce loss, 直接来看代码解析:
1 | # ../maskrcnn_benchmark/engine/trainer.py |
inference.py 文件概览
该文件定义了模型推演时的代码逻辑, 在拥有预训练的模型文件时, 可以通过调用本文件的函数来进行模型测试或模型推演, 文件的函数概览如下所示:
1 | # ./maskrcnn_benchmark/engine/inference.py |
inference 导入各种包及函数
我们首先来看看该文件导入了哪些包及函数
1 | # ./maskrcnn_benchmark/engine/inference.py |
evaluate
: 该函数封装了 coco_evaluation
和 voc_evaluation
, 具体解析可以查看data.
其余三个函数都来自 comm.py
文件, 详情着看 comm
inference.inference() 模型测试/推演核心代码
下面我们按照函数的调用顺序, 先来看看模型测试/推演的核心代码, 解析如下
1 | # ./maskrcnn_benchmark/engine/inference.py |
inference.compute_on_dataset() 计算结果
在上面的函数中, 调用了 compute_on_dataset()
函数来获得预测结果, 该函数的主要逻辑就是利用训练好的模型来预测模型的输出, 然后再将放在字典里返回, 代码解析如下:
1 | # ./maskrcnn_benchmark/engine/inference.py |
inference._accumulate_predictions_from_multiple_gpus()
由于 MaskrcnnBenchmark 默认是支持分布式的, 因此可以在多个 GPU 上计算结果, 但是最终需要把所有的结果都合并起来, 这正是 _accumulate_predictions_from_multiple_gpus()
的功能, 函数解析如下:
1 | # ./maskrcnn_benchmark/engine/inference.py |