pytorch减少显存使用
显存耗尽
RuntimeError: Caught RuntimeError in replica 0 on device 0.
RuntimeError: CUDA out of memory. Tried to allocate 4.01 GiB (GPU 0; 11.91 GiB total capacity; 4.89 GiB already allocated; 2.38 GiB free; 4.02 GiB cached)
transpose占用大量显存
通常transpose(),permute()会改变tensor的结构，需要使用.contiguous()将内存转换为连续的。估计pytorch为了加快计算速度，使用了占用大量显存的运算方式，开辟出新的显存并进行赋值。对这些函数进行修改可以减少大量显存：
将tensor分为多个batch，分别进行transpose
b,c,h,w = out.shape
# transpose dimension (1,2): bxcxhxw -> bxhxcxw. batchsize=h
# out = out.reshape(b, c, h, w).permute(0,2,1,3).contiguous()
batchsize = h
for ind_b in range(b):
for ind_r in range(0,w):
tmp = out[ind_b,:,:,ind_r*batchsize:(ind_r+1)*batchsize].transpose(0,1).contiguous()
out[ind_b,:,:,ind_r*batchsize:(ind_r+1)*batchsize] = tmp.reshape(N,c,batchsize)
矩阵点乘
矩阵点乘涉及到开辟新显存，如果用结果覆盖原有的数据可以节约显存：
# corr: b,c,h,w
# image_uf: b,c,h,w
# out = corr * image_uf
for batch in range(image_uf.shape[1]):
image_uf[:,batch,:,:] = image_uf[:,batch,:,:]*corr[:,batch,:,:]
注意，pytorch.sum()的结果与直接相加的结果不同：
# out = (corr * image_uf).sum(2)
for batch in range(image_uf.shape[1]):
image_uf[:,batch,:,:] = image_uf[:,batch,:,:]*corr[:,batch,:,:]
out = image_uf.sum(2)
不可以使用
# out = (corr * image_uf).sum(2)
out = torch.zeros(b,h,w)
for batch in range(image_uf.shape[1]):
out = out + image_uf[:,batch,:,:]*corr[:,batch,:,:]
注意
只能在测试代码中使用，要传梯度的话要再想办法计算梯度