如何在针对移动设备优化的模型上启用 Pytorch Dropout

Question

我有一个使用 dropout 进行 MC 预测的模型。它在我的桌面上的 Python 中运行良好。我想在移动平台上执行相同的算法，但 dropout 似乎没有应用。在移动设备上评估模型时，每次 MC 迭代都会得到相同的结果。

首先，我在 PyTorch 中加载模型并启用 dropout，如其他问题 [1] 中所述

model.eval()
for m in model.modules():
    if m.__class__.__name__.startswith('Dropout'):            
        m.train()

return model

接下来我针对移动设备进行优化。我添加了参数以禁用在优化方法期间删除丢失：

torchscript_model = torch.jit.script(model)
optimize_for_mobile(torchscript_model,
   optimization_blocklist={MobileOptimizerType.REMOVE_DROPOUT})
   ._save_for_lite_interpreter(ptlFile)

最后，在 Android 中，我加载模型然后执行 MC 迭代（简化示例）..

module = LiteModuleLoader.loadModuleFromAsset( this.getAssets(), "<ptlFile name>" );
<snip MC iterations>
float[] score = module.forward(IValue.from(inputTensor)).toTensor().getDataAsFloatArray();
</snip>

每次迭代我都会得到相同的分值。有什么建议如何在移动设备上实现与桌面设备相同的结果吗？如果我从优化器阻止列表中删除 REMOVE_DROPOUT，则导出的模型的大小会更小。因此，丢失层似乎正确地存在于移动模型中，但它们在移动评估期间似乎并未处于活动状态。

桌面版 Pytorch 版本 2.2.1。
Android 上的 pytorch_android_lite 1.13.1。

[1] 在 pytorch 上使用 MC Dropout 测量不确定性

Answer 1

尽管 REMOVE_DROPOUT 优化块正在努力将这些层保留在模型中，但似乎在针对移动设备进行优化期间，dropout 层正在转换回 eval 模式。解决方法是将 dropout 层替换为不考虑训练参数的自定义层：

class AlwaysDropout(torch.nn.Module):
    def __init__(self, dropout_prob):
        super(AlwaysDropout,self).__init__()
        self.dropout_prob = dropout_prob

    def forward(self,x):
        return torch.nn.functional.dropout2d(x,self.dropout_prob,True)

对于我的具体示例，我使用的是 monai 的密集网。为了交换 dropout 层，我只是替换了创建它们的工厂函数：

# overwrite the layer factory for monai Dropout
@monai.networks.layers.factories.Dropout.factory_function("dropout")
def dropout_factory(dim):
    return AlwaysDropout