作为机器学习从业者，咱们常常会遇到这样的状况，想要训练一个比拟大的模型，而 GPU 却因为内存不足而无奈训练它。当咱们在出于平安起因不容许在云计算的环境中工作时，这个问题常常会呈现。在这样的环境中，咱们无奈足够快地扩大或切换到功能强大的硬件并训练模型。并且因为梯度降落算法的性质，通常较大的批次在大多数模型中会产生更好的后果，但在大多数状况下，因为内存限度，咱们必须应用适应GPU显存的批次大小。

本文将介绍解梯度检查点（Gradient Checkpointing），这是一种能够让你以减少训练工夫为代价在 GPU 中训练大模型的技术。 咱们将在 PyTorch 中实现它并训练分类器模型。

梯度检查点

在反向流传算法中，梯度计算从损失函数开始，计算后更新模型权重。图中每一步计算的所有导数或梯度都会被存储，直到计算出最终的更新梯度。这样做会耗费大量 GPU 内存。梯度检查点通过在须要时从新计算这些值和抛弃在进一步计算中不须要的先前值来节俭内存。

让咱们用上面的虚构图来解释。

下面是一个计算图，每个叶节点上的数字相加失去最终输入。假如这个图示意反向流传期间产生的计算，那么每个节点的值都会被存储，这使得执行求和所需的总内存为7，因为有7个节点。然而咱们能够用更少的内存。假如咱们将1和2相加，并在下一个节点中将它们的值存储为3，而后删除这两个值。咱们能够对4和5做同样的操作，将9作为加法的后果存储。3和9也能够用同样的形式操作，存储后果后删除它们。通过执行这些操作，在计算过程中所需的内存从7缩小到3。

在没有梯度检查点的状况下，应用PyTorch训练分类模型

咱们将应用PyTorch构建一个分类模型，并在不应用梯度检查点的状况下训练它。记录模型的不同指标，如训练所用的工夫、内存耗费、准确性等。

因为咱们次要关注GPU的内存耗费，所以在训练时须要检测每批的内存耗费。这里应用nvidia-ml-py3库，该库应用nvidia-smi命令来获取内存信息。

 pip install nvidia-ml-py3

为了简略起见，咱们应用简略的狗和猫分类数据集的子集。

 git clone https://github.com/laxmimerit/dog-cat-full-dataset.git

执行上述命令后会在dog-cat-full-dataset的文件夹中失去残缺的数据集。

导入所需的包并初始化nvdia-smi

 importtorch importtorch.nnasnn importtorch.optimasoptim importnumpyasnp fromtorchvisionimportdatasets, models, transforms importmatplotlib.pyplotasplt importtime importos importcv2 importnvidia_smi importcopy fromPILimportImage fromtorch.utils.dataimportDataset,DataLoader importtorch.utils.checkpointascheckpoint fromtqdmimporttqdm importshutil fromtorch.utils.checkpointimportcheckpoint_sequential device="cuda"iftorch.cuda.is_available() else"cpu" %matplotlibinline importrandom  nvidia_smi.nvmlInit()

导入训练和测试模型所需的所有包。咱们还初始化nvidia-smi。

定义数据集和数据加载器

 #Define the dataset and the dataloader. train_dataset=datasets.ImageFolder(root="/content/dog-cat-full-dataset/data/train",                             transform=transforms.Compose([                                 transforms.RandomRotation(30),                                 transforms.RandomHorizontalFlip(),                                 transforms.RandomResizedCrop(224, scale=(0.96, 1.0), ratio=(0.95, 1.05)),                                 transforms.ToTensor(),                                 transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])                             ]))  val_dataset=datasets.ImageFolder(root="/content/dog-cat-full-dataset/data/test",                             transform=transforms.Compose([                                 transforms.Resize([224, 224]),                                 transforms.ToTensor(),                                 transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),                             ]))  train_dataloader=DataLoader(train_dataset,                             batch_size=64,                             shuffle=True,                             num_workers=2)  val_dataloader=DataLoader(val_dataset,                             batch_size=64,                             shuffle=True,                             num_workers=2)

这里咱们用torchvision数据集的ImageFolder类定义数据集。还在数据集上定义了某些转换，如RandomRotation, RandomHorizontalFlip等。最初对图片进行归一化，并且设置batch_size=64

定义训练和测试函数

 deftrain_model(model,loss_func,optimizer,train_dataloader,val_dataloader,epochs=10):      model.train()     #Training loop.     forepochinrange(epochs):       model.train()       forimages, targetintqdm(train_dataloader):           images, target=images.to(device), target.to(device)           images.requires_grad=True           optimizer.zero_grad()           output=model(images)           loss=loss_func(output, target)           loss.backward()           optimizer.step()       ifos.path.exists('grad_checkpoints/') isFalse:         os.mkdir('grad_checkpoints')       torch.save(model.state_dict(), 'grad_checkpoints/epoch_'+str(epoch)+'.pt')         #Test the model on validation data.       train_acc,train_loss=test_model(model,train_dataloader)       val_acc,val_loss=test_model(model,val_dataloader)        #Check memory usage.       handle=nvidia_smi.nvmlDeviceGetHandleByIndex(0)       info=nvidia_smi.nvmlDeviceGetMemoryInfo(handle)       memory_used=info.used       memory_used=(memory_used/1024)/1024        print(f"Epoch={epoch} Train Accuracy={train_acc} Train loss={train_loss} Validation accuracy={val_acc} Validation loss={val_loss} Memory used={memory_used} MB")    deftest_model(model,val_dataloader):   model.eval()   test_loss=0   correct=0   withtorch.no_grad():       forimages, targetinval_dataloader:           images, target=images.to(device), target.to(device)           output=model(images)           test_loss+=loss_func(output, target).data.item()           _, predicted=torch.max(output, 1)           correct+= (predicted==target).sum().item()      test_loss/=len(val_dataloader.dataset)    returnint(correct/len(val_dataloader.dataset) *100),test_loss

下面创立了一个简略的训练和测试循环来训练模型。最初还通过调用nvidia-smi计算内存应用。

训练

 torch.manual_seed(0)  #Learning rate. lr=0.003  #Defining the VGG16 sequential model. vgg16=models.vgg16() vgg_layers_list=list(vgg16.children())[:-1] vgg_layers_list.append(nn.Flatten()) vgg_layers_list.append(nn.Linear(25088,4096)) vgg_layers_list.append(nn.ReLU()) vgg_layers_list.append(nn.Dropout(0.5,inplace=False)) vgg_layers_list.append(nn.Linear(4096,4096)) vgg_layers_list.append(nn.ReLU()) vgg_layers_list.append(nn.Dropout(0.5,inplace=False)) vgg_layers_list.append(nn.Linear(4096,2)) model=nn.Sequential(*vgg_layers_list) model=model.to(device)    #Num of epochs to train num_epochs=10  #Loss loss_func=nn.CrossEntropyLoss()  # Optimizer  # optimizer = optim.Adam(model.parameters(), lr=lr, weight_decay=1e-5) optimizer=optim.SGD(params=model.parameters(), lr=0.001, momentum=0.9)   #Training the model. model=train_model(model, loss_func, optimizer,                        train_dataloader,val_dataloader,num_epochs)

咱们应用VGG16模型进行分类。上面是模型的训练日志。

能够从下面的日志中看到，在没有检查点的状况下，训练64个批大小的模型大概须要5分钟，占用内存为14222.125 mb。

应用带有梯度检查点的PyTorch训练分类模型

为了用梯度检查点训练模型，只须要编辑train_model函数。

 deftrain_with_grad_checkpointing(model,loss_func,optimizer,train_dataloader,val_dataloader,epochs=10):       #Training loop.     forepochinrange(epochs):       model.train()       forimages, targetintqdm(train_dataloader):           images, target=images.to(device), target.to(device)           images.requires_grad=True           optimizer.zero_grad()           #Applying gradient checkpointing           segments=2            # get the modules in the model. These modules should be in the order           # the model should be executed           modules= [modulefork, moduleinmodel._modules.items()]            # now call the checkpoint API and get the output           output=checkpoint_sequential(modules, segments, images)           loss=loss_func(output, target)           loss.backward()           optimizer.step()       ifos.path.exists('checkpoints/') isFalse:         os.mkdir('checkpoints')       torch.save(model.state_dict(), 'checkpoints/epoch_'+str(epoch)+'.pt')         #Test the model on validation data.       train_acc,train_loss=test_model(model,train_dataloader)       val_acc,val_loss=test_model(model,val_dataloader)        #Check memory.       handle=nvidia_smi.nvmlDeviceGetHandleByIndex(0)       info=nvidia_smi.nvmlDeviceGetMemoryInfo(handle)       memory_used=info.used       memory_used=(memory_used/1024)/1024        print(f"Epoch={epoch} Train Accuracy={train_acc} Train loss={train_loss} Validation accuracy={val_acc} Validation loss={val_loss} Memory used={memory_used} MB")  deftest_model(model,val_dataloader):   model.eval()   test_loss=0   correct=0   withtorch.no_grad():       forimages, targetinval_dataloader:           images, target=images.to(device), target.to(device)           output=model(images)           test_loss+=loss_func(output, target).data.item()           _, predicted=torch.max(output, 1)           correct+= (predicted==target).sum().item()      test_loss/=len(val_dataloader.dataset)    returnint(correct/len(val_dataloader.dataset) *100),test_lossdeftest_model(model,val_dataloader)

咱们将函数名批改为train_with_grad_checkpointing。也就是不通过模型(图)运行训练，而是应用checkpoint_sequential函数进行训练，该函数有三个输出:modules, segments, input。modules是神经网络层的列表，按它们执行的顺序排列。segments是在序列中创立的段的个数，应用梯度检查点进行训练以段为单位将输入用于从新计算反向流传期间的梯度。本文设置segments=2。input是模型的输出，在咱们的例子中是图像。这里的checkpoint_sequential仅用于程序模型，对于其余一些模型将产生谬误。

应用梯度检查点进行训练，如果你在notebook上执行所有的代码。倡议重新启动，因为nvidia-smi可能会取得以前代码中的内存耗费。

 torch.manual_seed(0)  lr=0.003  # model = models.resnet50() # model=model.to(device)  vgg16=models.vgg16() vgg_layers_list=list(vgg16.children())[:-1] vgg_layers_list.append(nn.Flatten()) vgg_layers_list.append(nn.Linear(25088,4096)) vgg_layers_list.append(nn.ReLU()) vgg_layers_list.append(nn.Dropout(0.5,inplace=False)) vgg_layers_list.append(nn.Linear(4096,4096)) vgg_layers_list.append(nn.ReLU()) vgg_layers_list.append(nn.Dropout(0.5,inplace=False)) vgg_layers_list.append(nn.Linear(4096,2)) model=nn.Sequential(*vgg_layers_list) model=model.to(device)     num_epochs=10  #Loss loss_func=nn.CrossEntropyLoss()  # Optimizer  # optimizer = optim.Adam(model.parameters(), lr=lr, weight_decay=1e-5) optimizer=optim.SGD(params=model.parameters(), lr=0.001, momentum=0.9)   #Fitting the model. model=train_with_grad_checkpointing(model, loss_func, optimizer,                        train_dataloader,val_dataloader,num_epochs)

输入如下：

从下面的输入能够看到，每个epoch的训练大概须要6分45秒。但只须要10550.125 mb的内存，也就是说咱们用工夫换取了空间，并且这两种状况下的精度都是79，因为在梯度检查点的状况下模型的精度没有损失。

总结

梯度检查点是一个十分好的技术，它能够帮忙在小显存的状况下残缺模型的训练。通过咱们的测试，个别状况下梯度检查点会将训练工夫缩短20%左右，然而工夫长点总比不能用要好，对吧。

本文的源代码：

https://avoid.overfit.cn/post/a13e29c312c741ac94d4a5079fb9f8af

作者：Vikas Kumar Ojha