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

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__()

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

39        self.act1 = nn.ReLU()

max-pooling

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

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

45        self.act2 = nn.ReLU()

max-pooling

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

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

51        self.act3 = nn.ReLU()

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):

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

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

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)

配置

79class Configs(MNISTConfigs):
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)

136        evidence = self.outputs_to_evidence(outputs)

139        loss = self.loss_func(evidence, target)

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

KL 背离损失系数

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, {
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()