使用PyTorch中的module.to（）移动成员张量

作者: 我头上有犄角
发布时间: 2024-12-11 05:29:02 (26天前)
转自：

3 条回复

0#
回复此人
雨儿 | 2019-08-31 10-32

<div class =“post-text”itemprop =“text”> <P> 用这个： </p> <pre> <code> device = torch.device("cuda" if torch.cuda.is_available() else "cpu") </code> </pre> <P> 现在您可以使用模型和每个张量 </p> <pre> <code> net.to(device) input = input.to(device) </code> </pre> </DIV>

编辑
1#
回复此人
永远的记忆 | 2019-08-31 10-32

<div class =“post-text”itemprop =“text”> <P> 经过一些试验和错误后，我发现了两种方法： </p> <OL> <LI> 使用缓冲区：通过替换 <code> self._train_noise = torch.randn(batch_size, embedding_size) </code> 同 <code> self.register_buffer('_train_noise', torch.randn(batch_size, embedding_size) </code> 噪声张量作为缓冲区添加到模块中。这让我们 <code> net.to(device) </code> 也影响它。此外，张量现在是state_dict的一部分。 </LI> <LI> <P> 覆盖 <code> net.to(device) </code> ：使用这个噪音不在state_dict之内。 </p> <pre> <code> def to(device): new_self = super(VariationalGenerator, self).to(device) new_self._train_noise = new_self._train_noise.to(device) new_self._eval_noise = new_self._eval_noise.to(device) return new_self </code> </pre> </LI> </醇> </DIV>

编辑

登录后才能参与评论