pytorch 就地操作显存节省

刺骨的言语ヽ痛彻心扉 2022-11-18 13:58 168阅读 0赞

X = torch.arange(12, dtype=torch.float32).reshape((3, 4))
    Y = torch.tensor([[2.0, 1, 4, 3], [1, 2, 3, 4], [4, 3, 2, 1]])

## 示例1 ##

before = id(Y)
    Y = Y + X
    id(Y) == before

以上代码输出结果为False。这表明进行矩阵加操作前的矩阵Y和进行操作后的矩阵Y本质上存储地址是不一样的。这主要是因为进行Y = Y + X操作时，pytorch会进行开辟一块新的存储空间，并分配新的地址，用于存储得到的结果，尽管二者的变量名是一样的。

但是实际使用中，进行矩阵加操作前的矩阵Y在得到结果后便不再需要了。如何使得进行上述操作后，存储结果存储到原本的矩阵Y的空间中去呢？

主要有两种写法

Y[:] = Y + X
    
    Y += X

## 示例2 ##

print('before id(Y):', id(Y))
    print(Y)
    Y[:] = Y + X
    print('after id(Y):', id(Y))
    print(Y)

以上代码输出结果为

before id(Y): 139876300140448
    tensor([[2., 1., 4., 3.],
            [1., 2., 3., 4.],
            [4., 3., 2., 1.]])
    after id(Y): 139876300140448
    tensor([[ 2.,  2.,  6.,  6.],
            [ 5.,  7.,  9., 11.],
            [12., 12., 12., 12.]])

可以看到Y的存储地址并没有改变，但内容改变了

## 示例3 ##

print('before id(Y):', id(Y))
    print(Y)
    Y += X
    print('after id(Y):', id(Y))
    print(Y)

输出结果为

before id(Y): 139876300138288
    tensor([[2., 1., 4., 3.],
            [1., 2., 3., 4.],
            [4., 3., 2., 1.]])
    after id(Y): 139876300138288
    tensor([[ 2.,  2.,  6.,  6.],
            [ 5.,  7.,  9., 11.],
            [12., 12., 12., 12.]])

同样，矩阵Y的地址并没有改变，内容改变了。

### 总结 ###

1. Y = Y + X 和 Y\[:\] = Y + X，Y += X 结果是一样的，但是所利用的存储空间大小不一样。

2. 如果后续代码不再使用变量X，建议使用就地操作减少存储空间（内存或者显存）的占用。