我正在 PyTorch 中构建变分自动编码器 (VAE)，但在编写与设备无关的代码时遇到了问题。 Autoencoder 是 nn.Module 的子项，具有编码器和解码器网络，它们也是。通过调用 net.to(device)，可以将网络的所有权重从一台设备转移到另一台设备。
我遇到的问题是重新参数化技巧:
encoding = mu + noise * sigma
噪声是一个与mu和sigma大小相同的张量，保存为autoencoder模块的成员变量。它在构造函数中初始化，并在每个训练步骤中就地重新采样。我这样做是为了避免每一步都构建一个新的噪声张量并将其推送到所需的设备。此外，我想修复评估中的噪音。这是代码:
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
当我现在使用 net.to('cuda:0') 将自动编码器移动到 GPU 时，我在转发时遇到错误，因为噪声张量没有移动。
我不想在构造函数中添加设备参数，因为那样以后仍然无法将其移动到另一个设备。我还尝试将噪声包装到 nn.Parameter 中，以便它受到 net.to() 的影响，但这会导致优化器出错，因为噪声是标记为 requires_grad=False。
谁有解决方案，可以使用 net.to() 移动所有模块？
请您参考如下方法:
tilman151's second approach 的更好版本可能是覆盖_apply，而不是to。这样 net.cuda()、net.float() 等都可以正常工作，因为它们都调用 _apply 而不是 to(如 the source 所示，比您想象的要简单):
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
