Pytorch释放显存占用方式

1. Pytorch释放显存占用方式

Pytorch是常用的机器学习框架之一，但是在训练深度学习模型时，经常会受到显存不足的问题的困扰。本文将介绍Pytorch释放显存占用的方式。

1.1 使用`torch.cuda.empty_cache()`释放显存空间

在Pytorch中，可以使用`torch.cuda.empty_cache()`手动释放显存。该函数会清空当前Pytorch分配的所有未释放显存空间。代码如下：

import torch
model = MyModel()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss()
for epoch in range(10):
    train_loss = 0.0
    model.train()
    
    for i, (inputs, labels) in enumerate(train_loader):
        inputs = inputs.cuda()
        labels = labels.cuda()
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        train_loss += loss.item()
    
    print('Epoch: {}, Train Loss: {}'.format(epoch+1, train_loss/len(train_loader)))
    torch.cuda.empty_cache() # 释放显存空间

然而，这种方式的效果并不是很明显。因此，需要使用其他更为有效的方式。

1.2 使用`torch.cuda.memory_allocated()`和`torch.cuda.memory_cached()`查看显存占用情况

在使用Pytorch训练深度学习模型时，经常需要查看显存的占用情况，以便及时释放占用的显存空间。可以使用`torch.cuda.memory_allocated()`和`torch.cuda.memory_cached()`函数查看当前Pytorch分配的显存占用情况，代码如下：

import torch
model = MyModel()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss()
print('Before Training: {} MB'.format(torch.cuda.memory_allocated()/1024/1024))
print('Before Training: {} MB'.format(torch.cuda.memory_cached()/1024/1024))
for epoch in range(10):
    train_loss = 0.0
    model.train()
    
    for i, (inputs, labels) in enumerate(train_loader):
        inputs = inputs.cuda()
        labels = labels.cuda()
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        train_loss += loss.item()
    
    print('Epoch: {}, Train Loss: {}'.format(epoch+1, train_loss/len(train_loader)))
    print('During Training: {} MB'.format(torch.cuda.memory_allocated()/1024/1024))
    print('During Training: {} MB'.format(torch.cuda.memory_cached()/1024/1024))
    torch.cuda.empty_cache() # 释放显存空间
print('After Training: {} MB'.format(torch.cuda.memory_allocated()/1024/1024))
print('After Training: {} MB'.format(torch.cuda.memory_cached()/1024/1024))

1.3 使用`with torch.cuda.device()`设置显存设备

可以使用`with torch.cuda.device()`函数手动设置使用哪个显卡设备，并释放掉没有使用的显存空间，以提高显存利用效率。代码如下：

import torch
device = torch.device('cuda:0')
model = MyModel()
model.to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss()
print('Before Training: {} MB'.format(torch.cuda.memory_allocated(device)/1024/1024))
print('Before Training: {} MB'.format(torch.cuda.memory_cached(device)/1024/1024))
for epoch in range(10):
    train_loss = 0.0
    model.train()
    
    with torch.cuda.device(device):
        for i, (inputs, labels) in enumerate(train_loader):
            inputs = inputs.to(device)
            labels = labels.to(device)
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            train_loss += loss.item()
    
    print('Epoch: {}, Train Loss: {}'.format(epoch+1, train_loss/len(train_loader)))
    print('During Training: {} MB'.format(torch.cuda.memory_allocated(device)/1024/1024))
    print('During Training: {} MB'.format(torch.cuda.memory_cached(device)/1024/1024))
    torch.cuda.empty_cache() # 释放显存空间
print('After Training: {} MB'.format(torch.cuda.memory_allocated(device)/1024/1024))
print('After Training: {} MB'.format(torch.cuda.memory_cached(device)/1024/1024))

需要注意的是，在`with torch.cuda.device()`中设置的显存设备，一定要与`model.to(device)`中设置的设备一致，否则会报错。

综上所述，释放显存占用的方式主要有以下三种：手动释放显存空间、查看显存占用情况、手动设置使用哪个显卡设备。需要根据实际情况选择合适的释放方式。