#

基于证据的深度学习量化分类不确定性实验

14from typing import Any
15
16import torch.nn as nn
17import torch.utils.data
18
19from labml import tracker, experiment
20from labml.configs import option, calculate
21from labml_helpers.module import Module
22from labml_helpers.schedule import Schedule, RelativePiecewise
23from labml_helpers.train_valid import BatchIndex
24from labml_nn.experiments.mnist import MNISTConfigs
25from labml_nn.uncertainty.evidence import KLDivergenceLoss, TrackStatistics, MaximumLikelihoodLoss, \
26    CrossEntropyBayesRisk, SquaredErrorBayesRisk

#

基于 LeNET 的 MINST 分类模型

29class Model(Module):

#

34    def __init__(self, dropout: float):
35        super().__init__()

#

第一个 $5 x 5$ 卷积层

37        self.conv1 = nn.Conv2d(1, 20, kernel_size=5)

#

激活 ReLU

39        self.act1 = nn.ReLU()

#

$2 x 2$ max-pooling

41        self.max_pool1 = nn.MaxPool2d(2, 2)

#

第二个 $5 x 5$ 卷积层

43        self.conv2 = nn.Conv2d(20, 50, kernel_size=5)

#

激活 ReLU

45        self.act2 = nn.ReLU()

#

$2 x 2$ max-pooling

47        self.max_pool2 = nn.MaxPool2d(2, 2)

#

第一个映射到要素的完全连接的 $500$ 图层

49        self.fc1 = nn.Linear(50 * 4 * 4, 500)

#

激活 ReLU

51        self.act3 = nn.ReLU()

#

最后一个完全连接的层，用于输出 $10$ 课堂证据。RelU 或 Softplus 激活在模型之外应用于此，以获得非负面证据

55        self.fc2 = nn.Linear(500, 10)

#

隐藏图层的退出

57        self.dropout = nn.Dropout(p=dropout)

#

x 是一批 MNIST 形状的图像[batch_size, 1, 28, 28]

59    def __call__(self, x: torch.Tensor):

#

应用第一个卷积和最大池。结果有形状[batch_size, 20, 12, 12]

65        x = self.max_pool1(self.act1(self.conv1(x)))

#

应用第二个卷积和最大池。结果有形状[batch_size, 50, 4, 4]

68        x = self.max_pool2(self.act2(self.conv2(x)))

#

将张量展平成形状[batch_size, 50 * 4 * 4]

70        x = x.view(x.shape[0], -1)

#

应用隐藏层

72        x = self.act3(self.fc1(x))

#

申请退学

74        x = self.dropout(x)

#

应用最后一层然后返回

76        return self.fc2(x)

#

配置

我们使用MNISTConfigs 配置。

79class Configs(MNISTConfigs):

#

KL 分歧正则化

87    kl_div_loss = KLDivergenceLoss()

#

KL 发散正则化系数时间表

89    kl_div_coef: Schedule

#

KL 发散正则化系数时间表

91    kl_div_coef_schedule = [(0, 0.), (0.2, 0.01), (1, 1.)]

#

用于跟踪的统计模块

93    stats = TrackStatistics()

#

辍学

95    dropout: float = 0.5

#

用于将模型输出转换为非零证据的模块

97    outputs_to_evidence: Module

#

初始化

99    def init(self):

#

设置跟踪器配置

104        tracker.set_scalar("loss.*", True)
105        tracker.set_scalar("accuracy.*", True)
106        tracker.set_histogram('u.*', True)
107        tracker.set_histogram('prob.*', False)
108        tracker.set_scalar('annealing_coef.*', False)
109        tracker.set_scalar('kl_div_loss.*', False)

#

112        self.state_modules = []

#

培训或验证步骤

114    def step(self, batch: Any, batch_idx: BatchIndex):

#

训练/评估模式

120        self.model.train(self.mode.is_train)

#

将数据移动到设备

123        data, target = batch[0].to(self.device), batch[1].to(self.device)

#

一热编码目标

126        eye = torch.eye(10).to(torch.float).to(self.device)
127        target = eye[target]

#

在训练模式下更新全局步长（处理的样本数）

130        if self.mode.is_train:
131            tracker.add_global_step(len(data))

#

获取模型输出

134        outputs = self.model(data)

#

获取证据 $e_{k} \geq 0$

136        evidence = self.outputs_to_evidence(outputs)

#

计算损失

139        loss = self.loss_func(evidence, target)

#

计算 KL 背离正则化损失

141        kl_div_loss = self.kl_div_loss(evidence, target)
142        tracker.add("loss.", loss)
143        tracker.add("kl_div_loss.", kl_div_loss)

#

KL 背离损失系数 $λ_{t}$

146        annealing_coef = min(1., self.kl_div_coef(tracker.get_global_step()))
147        tracker.add("annealing_coef.", annealing_coef)

#

总亏损

150        loss = loss + annealing_coef * kl_div_loss

#

追踪统计数据

153        self.stats(evidence, target)

#

训练模型

156        if self.mode.is_train:

#

计算梯度

158            loss.backward()

#

采取优化器步骤

160            self.optimizer.step()

#

清除渐变

162            self.optimizer.zero_grad()

#

保存跟踪的指标

165        tracker.save()

#

创建模型

168@option(Configs.model)
169def mnist_model(c: Configs):

#

173    return Model(c.dropout).to(c.device)

#

KL 背离损失系数时间表

176@option(Configs.kl_div_coef)
177def kl_div_coef(c: Configs):

#

创建相对的分段时间表

183    return RelativePiecewise(c.kl_div_coef_schedule, c.epochs * len(c.train_dataset))

#

最大似然损失

187calculate(Configs.loss_func, 'max_likelihood_loss', lambda: MaximumLikelihoodLoss())

#

交叉熵贝叶斯风险

189calculate(Configs.loss_func, 'cross_entropy_bayes_risk', lambda: CrossEntropyBayesRisk())

#

平方误差贝叶斯风险

191calculate(Configs.loss_func, 'squared_error_bayes_risk', lambda: SquaredErrorBayesRisk())

#

RelU 来计算证据

194calculate(Configs.outputs_to_evidence, 'relu', lambda: nn.ReLU())

#

Softplus 用于计算证据

196calculate(Configs.outputs_to_evidence, 'softplus', lambda: nn.Softplus())

#

199def main():

#

创建实验

201    experiment.create(name='evidence_mnist')

#

创建配置

203    conf = Configs()

#

装载配置

205    experiment.configs(conf, {
206        'optimizer.optimizer': 'Adam',
207        'optimizer.learning_rate': 0.001,
208        'optimizer.weight_decay': 0.005,

#

'loss_func'：'max_imilihood_loss'，'loss_func'：'cross_entropy_bayes_risk '，

212        'loss_func': 'squared_error_bayes_risk',
213
214        'outputs_to_evidence': 'softplus',
215
216        'dropout': 0.5,
217    })

#

开始实验并运行训练循环

219    with experiment.start():
220        conf.run()

#

224if __name__ == '__main__':
225    main()