CompressAI 模型参数配置学习笔记

CompressAI 自己训练网络的参数设置

• -m

模型设置，通过 -m 选择要用来训练的模型，默认为 "bmshj2018_factorized"，CompressAI 包含的所有模型如下列表所示：

__all__ = [
    "bmshj2018_factorized",
    "bmshj2018_factorized_relu",
    "bmshj2018_hyperprior",
    "mbt2018",
    "mbt2018_mean",
    "cheng2020_anchor",
    "cheng2020_attn",
    "pre_mbt2018"           # 自己加的结合前处理网络的 mbt2018
]

• –epochs

时期，表示在训练集上训练多少次，一个Epoch就是将所有训练样本训练一次的过程。官方默认 100

• –lambda

该表格为 CompressAI 在训练端到端编码器的时候使用的 lambda 值，作用相当于大致确定一个码率点，默认值为1e-2

• –learning-rate

学习率，越大学习梯度越快，默认为 1e-4

• –batch-size

每一次训练选取的输入数量，CompressAI 官方训练端到端 Codec 的值为 16 / 32

• –test-batch-size

同 --batch-size 的作用，不过此处是训练输入的选取，官方默认值为 64

• –patch-size

补丁大小，意为每次操作的时候选取的块的大小，官方默认值为 (256, 256)

• –cuda

开关打开后，使用 cuda 进行训练，即 GPU

• –save

开关打开后，会实时保存训练好的模型，保存的地址为 examples\checkpoint.pth.tar 以及 examples\checkpoint_best_loss.pth.tar，一个是每次训练都会保存，一个是只有loss小于之前的最小值才会保存。

• Loss-Function

训练时使用的损失函数表达式如下所示：