关于人工智能:在-PyTorch-中使用梯度检查点在GPU-上训练更大的模型

38次阅读

共计 7917 个字符，预计需要花费 20 分钟才能阅读完成。

作为机器学习从业者，咱们常常会遇到这样的状况，想要训练一个比拟大的模型，而 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

正文完
人工智能
发表至： 人工智能
2023-02-01
 0
关于人工智能:AI技术在小程序生态质量保障方向的落地实践
关于人工智能:QA-特辑-百万员工不上班也能领全勤奖的原因都在这里了
关于人工智能:究诸经典探寻大模型演变之踪迹
关于人工智能:OpenAIs-services-are-not-available-in-your-country