在 PyTorch 中使用 module.to() 移动成员张量

我正在 PyTorch 中构建变分自动编码器 (VAE)，但在编写与设备无关的代码时遇到问题。自动编码器是nn.Module具有编码器和解码器网络，它们也是。网络的所有权重都可以通过调用从一台设备移动到另一台设备net.to(device).

我遇到的问题是重新参数化技巧：

encoding = mu + noise * sigma

噪声是一个与以下大小相同的张量mu and sigma并保存为自动编码器模块的成员变量。它在构造函数中初始化，并在每个训练步骤中就地重新采样。我这样做是为了避免每一步构建一个新的噪声张量并将其推送到所需的设备。此外，我想修复评估中的噪音。这是代码：

class VariationalGenerator(nn.Module):
    def __init__(self, input_nc, output_nc):
        super(VariationalGenerator, self).__init__()

        self.input_nc = input_nc
        self.output_nc = output_nc
        embedding_size = 128

        self._train_noise = torch.randn(batch_size, embedding_size)
        self._eval_noise = torch.randn(1, embedding_size)
        self.noise = self._train_noise

        # Create encoder
        self.encoder = Encoder(input_nc, embedding_size)
        # Create decoder
        self.decoder = Decoder(output_nc, embedding_size)

    def train(self, mode=True):
        super(VariationalGenerator, self).train(mode)
        self.noise = self._train_noise

    def eval(self):
        super(VariationalGenerator, self).eval()
        self.noise = self._eval_noise

    def forward(self, inputs):
        # Calculate parameters of embedding space
        mu, log_sigma = self.encoder.forward(inputs)
        # Resample noise if training
        if self.training:
            self.noise.normal_()
        # Reparametrize noise to embedding space
        inputs = mu + self.noise * torch.exp(0.5 * log_sigma)
        # Decode to image
        inputs = self.decoder(inputs)

        return inputs, mu, log_sigma

当我现在将自动编码器移动到 GPU 时net.to('cuda:0')我在转发时遇到错误，因为噪声张量没有移动。

我不想向构造函数添加设备参数，因为这样以后仍然无法将其移动到另一个设备。我也尝试将噪音包裹起来nn.Parameter从而使其受到影响net.to()，但这会给优化器带来错误，因为噪声被标记为requires_grad=False.

任何人都有一个解决方案来移动所有模块net.to()?

更好的版本tilman151的第二种方法 https://stackoverflow.com/a/54768936/344821可能是覆盖_apply， 而不是to。那样net.cuda(), net.float()等都可以工作，因为它们都调用_apply而不是to（可以看出来源 https://github.com/pytorch/pytorch/blob/master/torch/nn/modules/module.py，这比你想象的要简单）：

def _apply(self, fn):
    super(VariationalGenerator, self)._apply(fn)
    self._train_noise = fn(self._train_noise)
    self._eval_noise = fn(self._eval_noise)
    return self