本文所述的内存优化方法不包含例如“减少batch size”等直接影响训练流的方法
目录
主要参考资料:部分其他资料将在文中直接给出
1.尽量使用in_place实现
使用in_place操作使得Pytorch的allocator不会记录该部分的原tensor,从而减少显存的消耗。也正是因为如此,如果在网络反向计算梯度的过程中需要用到原tensor,那么就不能够使用in_place操作,否则会使得反向传播报错。
ReLU:
Pytorch的nn.ReLU支持in_place操作,一般原则是,只要ReLU的前一层“不需要”进行梯度计算,那么ReLU就可以使用in_place实现:例如,当ReLU层前是一个BN层时(Pytorch的BN会在反向传播过程中重新计算对应的需求项,因此无需记录梯度信息)。
2.Pytorch的 ‘checkpoint’
torch.utils.checkpoint提供了一种用 “时间换空间” 的解决思路,即对于部分function
或者model
,我们可以用checkpoint
对其进行包装,被包装的model,在正向传播的过程中不会记录对应的结果(节省空间),而是在反向传播的过程中重新计算(增加了计算成本)。
除此以外,使用checkpoint还有很多需要注意的细节,如果使用不当可能造成“两次”计算的结果是不相等的,这样就产生了误差,是我们不想看到的,更多细节可以参考 官方文档 。
小案例:
model = nn.Sequential(nn.Conv2d(3,6,1,1), nn.ReLU()).to('cuda')
dummy_inp = torch.randn(1,3,128,128).to('cuda')
dummy_inp = autograd.Variable(dummy_input, requires_grad=True)
out = checkpoint_sequential(model, 2, dummy_input)
# 注意,如果这里的dummy_input没有requires_grad=True,会报错,因为checkpoint一定是在需要求梯度的情况下使用的,否则本来就不需要保存原值。
这里在推荐两个基于Pytorch checkpoint特性实现的giuhub库,实现了诸如将 batchnorm 和relu打包成inplace操作等:
https://github.com/gpleiss/efficient_densenet_pytorch
https://github.com/mapillary/inplace_abn
3.删除loss等不再被需要的值,释放对应的空间
根据大家的说法,这个方法其实作用并不是很大,但是我们姑且还是“了解一下”,其实它的思想很简单,一次循环之后如果我们已经完成了对应的optim以及loss追溯的操作,那么output, loss等值当然就不再被需要了,这时候手动释放这些内存可能是值得尝试的。
for i, (x, y) in enumerate(train_loader):
x = Variable(x)
y = Variable(y)
# compute model and update
del x, y, output
注意:如果数据集很小,这个操作可能得不偿失
4.混合精度训练
一般而言,Pytorch的模型训练要求数据为float32的精度,但事实上不同的训练子层对数据精度的需求是不一致的,完全采用float32的数据精度会造成很大的存储空间浪费,这也为空间优化提供了一个方向。目前混合精度训练(AMP)即有三方的实现也有Pytorch的官方集成版本:
这里对官方提供的方法进行简单的介绍:(更详细的内容请参考官方文档)
- torch.cuda.amp.autocast(enabled=True)
autocast有上下文管理器(context manager)以及装饰器(decorator)两种使用方法
- context manager
# 定义模型以及优化器
model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)
for input, target in data:
optimizer.zero_grad()
# 上下文管理器启动
with autocast():
output = model(input)
loss = loss_fn(output, target)
# 在backward阶段退出上下文管理器
loss.backward()
optimizer.step()
注意:文档中推荐在完成forward后退出上下文管理器,因为如果在BP的过程中仍然处于管理器中,那么在自动调节的机制下可能会出现类型不匹配错误,如下所述:
- decorator
class AutocastModel(nn.Module):
...
@autocast()
def forward(self, input):
...
作为装饰器的使用方法十分简单,只需要在对应模型的FW的方法前进行装饰即可。
- 手动恢复数据类型,避免数据不匹配错误
在我们对模型的精度需求很明确的情况下,为了避免退出管理器后出现精度匹配的错误,我们可以进行手动地恢复,下面再挂一段官方案例:
a_float32 = torch.rand((8, 8), device="cuda")
b_float32 = torch.rand((8, 8), device="cuda")
c_float32 = torch.rand((8, 8), device="cuda")
d_float32 = torch.rand((8, 8), device="cuda")
with autocast():
# 假设torch.mm需求的是torch16精度的数据输入,输出同样为F16
# 由于管理器的存在,我们不需要手动进行数据类型转换
e_float16 = torch.mm(a_float32, b_float32)
f_float16 = torch.mm(d_float32, e_float16)
# 这里已经退出了上下文管理器,所以需要手动进行类型转换
# 注意,虽然torch.mm运行在float16下,但是输入32的数据精度是没有问题的
g_float32 = torch.mm(d_float32, f_float16.float())
文档中的补充说明:在autocast管理器中我们是无需担心数据不匹配错误的,但并不一定不存在,也要注意。
- autocast(enable=False)嵌套在autocast()中使用
这种方式是受推荐的,一些情况下我们可能想要特定的阶段运行在特定的数据类型精度下,那么只需要适时地关闭上下文管理器即可:
a_float32 = torch.rand((8, 8), device="cuda")
b_float32 = torch.rand((8, 8), device="cuda")
c_float32 = torch.rand((8, 8), device="cuda")
d_float32 = torch.rand((8, 8), device="cuda")
with autocast():
e_float16 = torch.mm(a_float32, b_float32)
with autocast(enabled=False):
# 保证运行在float32的精度下
f_float32 = torch.mm(c_float32, e_float16.float())
g_float16 = torch.mm(d_float32, f_float32)
- 自定义装饰器,指定数据类型
使用torch.cuda.amp.custom_fwd(fwd=None, **kwargs)
以及torch.cuda.amp.custom_bwd(fwd=None, **kwargs)
分别针对forward以及backward pass,能够实现自定义的数据精度。使用方法就是装饰器。需要注意的是:这两个装饰器只有当位于autocast()上下文管理器中才能够发挥作用,否则不产生任何行为。
- torch.cuda.amp.GradScaler
在float16精度下,当gradient的量级过小时可能会被忽略(归零)“下溢”,这对训练是十分不利的。GradScaler的思想是十分朴素的,即在进行backward之前,先对放大(scale)loss的值,直至避免或减少下溢发生。官方描述如下:
定义方法:
GradScaler(init_scale=65536.0, growth_factor=2.0, backoff_factor=0.5, growth_interval=2000, enabled=True)
使用方法:
# 训练阶段开始前创建GradScalar
| scaler = GradScaler()
|
| for epoch in epochs:
| for input, target in data:
| optimizer.zero_grad()
| output = model(input)
| loss = loss_fn(output, target)
|
| # 先对loss进行放大,在backward
| scaler.scale(loss).backward()
|
# 用sclar包装optimizer
# scaler首先unscale各parameter的gradient值
# 注意:如果缩小后的梯度值中没有产生上溢(inf/NaNs),那么接着调用optimizer.step(),否则跳过optimizer.step()
| scaler.step(optimizer)
|
| # 依据定义的参数,更新scale值
# 这里的更新依据是在放缩后是否出现了上溢
# 如果有上溢那么backoff_factor起作用
# 如果没有上溢那么在growth_interval个轮次之后,growth_factor起作用。
| scaler.update()
注意:如果enabled=False那么scalar不会起作用,而是直接调用optimizer.step()。
5.model.eval()以及with torch.no_grad()
除了train阶段,在无需进行bcw的地方,不仅仅是从节省显存的角度考虑,也有助于我们更好地把控整个网络:如在validation阶段我们可以通过上下文管理器model.eval(),将各模型调整至evaluation阶段;或者直接通过with torch.no_grad()避免bwd行为。需要注意的是:model.eval()与with torch.no_grad()并不是等价的。具体而言:model.eval()会改变各网络层次的行为,如Dropout层,在train阶段与eval阶段的行为表现存在差异。可以参考:https://discuss.pytorch.org/t/model-eval-vs-with-torch-no-grad/19615
6.torch.cuda.empty_cache()
该命令的作用可以简述为:释放Pytorch的memory allocator所占用的“无用”的空间。特别需要注意的是“无用的”三个字,也就是说该命令是无法直接根据我们的需要释放指定的空间的,而由分配器统一管理。大多数情况下该命令的效果并不显著,但也可以试试
。
7.分割训练过程(基于checkpoint)
假设网络层次太深,那么可以将网络进行切割,先对前部分进行训练再对后部分进行训练。例如网络有100层,一次性训练时显存无法支持,那么就进行切割,将其分为网络1与网络2,分别训练。这里贴一段示例代码:注:代码来自 https://oldpan.me/archives/how-to-use-memory-pytorch
这里需要用到我们前面介绍的checkpoint特性, 机制我们在前面也已经进行介绍了
input = torch.rand(1, 10)
layers = [nn.Linear(10, 10) for _ in range(1000)]
model = nn.Sequential(*layers)
# 方式一:
# output = model(input)
# 进行如下更改
# 设置输入的input=>requires_grad=True
# 如果不设置可能会导致得到的gradient为0
input = torch.rand(1, 10, requires_grad=True)
layers = [nn.Linear(10, 10) for _ in range(1000)]
# 定义要计算的层函数
# 一个计算前500个层,另一个计算后500个层
def run_first_half(*args):
x = args[0]
for layer in layers[:500]:
x = layer(x)
return x
def run_second_half(*args):
x = args[0]
for layer in layers[500:-1]:
x = layer(x)
return x
# 引入新加的checkpoint
from torch.utils.checkpoint import checkpoint
x = checkpoint(run_first_half, input)
x = checkpoint(run_second_half, x)
# 最后一层单独调出来执行
x = layers[-1](x)
x.sum.backward()
对于nn.Sequiential的情况则可以:
input = torch.rand(1, 10, requires_grad=True)
layers = [nn.Linear(10, 10) for _ in range(1000)]
model = nn.Sequential(*layers)
from torch.utils.checkpoint import checkpoint_sequential
# 分成两个部分
num_segments = 2
x = checkpoint_sequential(model, num_segments, input)
x.sum().backward()
注意:该方法会影响与batchsize相关的层,例如batchNorm层
这里补充一段官方文档的注释:
从上图中可以看出:最后一个segment不会运行在‘torch.no_grad()’模式下,也就是说如果我们将segments设置为1,那么checkpoint_sequential将表现为无效。
8.分批训练,一次backward(loss accumulation)
即减小batchsize的另一种形式,假设当batchsize=64时显存溢出,那么我们可以设定batchsize为32,训练两次,累计误差再一次性backward。需要注意:该方法会影响与batchsize大小相关的层次,例如BatchNorm层
。示例:
# loss accumulation
accumulation_step = 4
for i in range(K):
train_dataset = tracking_dataset(kitti_object, root_dir=train_root, ki=i, K=K, typ='train')
val_dataset = tracking_dataset(kitti_object, root_dir=train_root, ki=i, K=K, typ='val')
train_loader = DataLoader(train_dataset, batch_size=2, drop_last=True)
val_loader = DataLoader(val_dataset, batch_size=2, drop_last=True)
for epoch in range(50):
for idx, batch in enumerate(train_loader):
print("batch{}.....".format(idx))
for key in batch:
batch[key] = batch[key].to(device)
inp = batch['inp']
# output
with autocast(enabled=True):
output = model(inp)
# loss
total_loss, losses = lossModel(output, batch)
loss_mean = total_loss.mean()
loss_mean /= accumulation_step
# accumulate
loss_mean.backward()
if (idx+1) % accumulation_step == 0:
# 一定次数后,计算梯度
optimizer.zero_grad()
optimizer.step()
print('success !!')