深度学习笔记--解决GPU显存使用量不断增加的问题

目录

1--问题描述

2--问题解决

3--代码

---

1--问题描述

        基于 Pytorch 使用 VGG16 预训练模型进行分类预测时，出现 GPU 显存使用量不断增加，最终出现 cuda out of memory 的问题；

        出现上述问题的原因在于：输入数据到网络模型进行推理时，会默认构建计算图，便于后续反向传播进行梯度计算。而构建完整的计算图，会增加计算和累积内存消耗，从而导致 GPU显存使用量不断增加；

        由于博主只使用 VGG16 预训练模型进行分类预测，不需要训练和反向传播更新参数，所以不用构建完整的计算图。

2--问题解决

        在推理代码中增加以下指令，表明当前计算不需要进行反向传播，即强制不进行完整计算图的构建：

with torch.no_grad():
    ...
    ...

3--代码

        问题代码：

def extract_rgb_feature(rgb_data):
    data = rgb_data.to(device_id[0]) # [40, 40, 3]
    data = data.permute(2, 0, 1).unsqueeze(0) # [1, 3, 40, 40]
    data = F.interpolate(data, size = (224, 224), mode='nearest').float() #[1, 3, 224, 224]
    data = model(data) # [1, linear_Class]
    return data

        修正代码：

def extract_rgb_feature(rgb_data):
    with torch.no_grad():
        data = rgb_data.to(device_id[0]) # [40, 40, 3]
        data = data.permute(2, 0, 1).unsqueeze(0) # [1, 3, 40, 40]
        data = F.interpolate(data, size = (224, 224), mode='nearest').float() #[1, 3, 224, 224]
        data = model(data) # [1, linear_Class]
        return data

        完整代码：

from torchvision import models
import torch.nn as nn
import torch
import numpy as np
import cv2
import torch.nn.functional as F

class My_Net(nn.Module):
    def __init__(self, linear_Class):
        super(My_Net, self).__init__()
        self.linear_Class = linear_Class
        self.backbone = models.vgg16(pretrained=True) # 以 vgg16 作为 backbone
        self.backbone = self.process_backbone(self.backbone) # 对预训练模型进行处理
 
        self.linear1 = nn.Linear(in_features = 4096, out_features = self.linear_Class)
 
    def process_backbone(self, model):
 
        # 固定预训练模型的参数
        for param in model.parameters():
            param.requires_grad = False
        
        # 删除最后预测层    
        del model.classifier[6]
 
        return model
    
    def forward(self, x):
        x = self.backbone(x)
        x = self.linear1(x)
        return x

linear_Class = 2
device_id = [7]
model = My_Net(linear_Class).to(device_id[0]) # 初始化模型

def extract_rgb_feature(rgb_data):
    with torch.no_grad():
        data = rgb_data.to(device_id[0]) # [40, 40, 3]
        data = data.permute(2, 0, 1).unsqueeze(0) # [1, 3, 40, 40]
        data = F.interpolate(data, size = (224, 224), mode='nearest').float() #[1, 3, 224, 224]
        data = model(data) # [1, linear_Class]
        return data
                    
     
if __name__ == "__main__":

    CSub_train_txt_path = '../statistics/CSub_train.txt'
    CSub_test_txt_path = '../statistics/CSub_test.txt'
    
    CSub_train_data_path = './2J_rgb_patch_npy_file_40x40/CSub/train/'
    CSub_test_data_path = './2J_rgb_patch_npy_file_40x40/CSub/test/'
    
    CSub_train_txt = np.loadtxt(CSub_train_txt_path, dtype = str)
    CSub_test_txt = np.loadtxt(CSub_test_txt_path, dtype = str)
    
    CSub_train_save_path = './pre_vgg_feature/2J/CSub/train.npy'
    CSub_test_save_path = './pre_vgg_feature/2J/CSub/test.npy'
    
    save_data = []
    
    for (idx, name) in enumerate(CSub_test_txt):
        data_path = CSub_test_data_path + name + '.npy' 
        rgb_data = np.load(data_path) # T, M, N, H, W, C
        rgb_data = torch.from_numpy(rgb_data)#.to(device = device_id[0])
        
        T, M, N, H, W, C = rgb_data.shape
        Output = torch.zeros(T, M, N, 1, linear_Class)
        
        for t in range(T):
            for m in range(M):
                for n in range(N):
                    data = extract_rgb_feature(rgb_data[t, m, n])
                    Output[t, m, n] = data.cpu()
                    
        save_data.append(Output) 
        print("Processing " + name + ", Done !")
        
    np.save(CSub_test_save_path, save_data)
        
    print("All done!")