调用使用GPU已有张量的方法时GPU内存占用异常增加的原因排查

代码背景 #

我有两段PyTorch代码，除了代码2引入了外部模块的extra_class外，DataLoader及其他所有逻辑完全一致：

代码1 #

main(args)
......
train(nn)
......

代码2 #

main(args)
......
object = extra_class(...arguments...)       # 创建对象
train(nn)
object.method1()                           # 调用对象方法
......

外部模块的extra_class实现 #

class extra_class(...parameters......):
    def __init__(...):
         tensor1.to(device)  # 对大量张量使用`.to(device)`
         tensor2.to(device) 
         ..................
    def method1(self): 
        ....  # 使用__init__()中已转移至设备的张量

问题现象 #

运行代码1时，GPU内存占用稳定（例如维持在70%）；但运行代码2时，执行object.method1()期间GPU内存突然飙升，执行完该行后内存回落。可推断该行执行时大量张量被加载到GPU内存，导致执行时间增加。但我无法理解原因：method1中并未使用.to(device)，仅使用__init__()中已转移至GPU设备的张量，按道理所有张量转移应在对象创建阶段完成，为何调用method1时才发生内存增长？

原因分析 #

问题出在__init__里的.to(device)调用逻辑：
PyTorch的.to()方法默认返回一个新的张量实例，而非原地修改原张量。你当前的代码只是调用了tensor1.to(device)，但没有把返回的GPU张量赋值给任何实例变量，原tensor1、tensor2仍然停留在CPU上，__init__里的转移操作等于没有生效。

等到method1执行时，代码使用这些CPU张量参与运算时，PyTorch会自动将CPU张量隐式转移到GPU（因为运算所在设备是GPU），这才导致了GPU内存的突然飙升。

解决方法 #

在__init__中必须将.to(device)的返回值赋值给类的实例变量，确保张量真正被转移到GPU：

class extra_class(...parameters......):
    def __init__(...):
         self.tensor1 = tensor1.to(device)  # 赋值给实例变量
         self.tensor2 = tensor2.to(device) 
         ..................
    def method1(self): 
        # 使用self.tensor1、self.tensor2进行运算
        ....

这样在对象创建阶段，张量就已经完成了GPU转移，后续method1执行时不会再触发隐式的设备转移，GPU内存也就不会出现突增的情况。

内容的提问来源于stack exchange，提问作者hanugm