目錄

1--問題描述

2--問題解決

3--代碼

1--問題描述

??????? 基于 Pytorch 使用 VGG16 預(yù)訓(xùn)練模型進(jìn)行分類預(yù)測時，出現(xiàn) GPU 顯存使用量不斷增加，最終出現(xiàn) cuda out of memory 的問題；

??????? 出現(xiàn)上述問題的原因在于：輸入數(shù)據(jù)到網(wǎng)絡(luò)模型進(jìn)行推理時，會默認(rèn)構(gòu)建計(jì)算圖，便于后續(xù)反向傳播進(jìn)行梯度計(jì)算。而構(gòu)建完整的計(jì)算圖，會增加計(jì)算和累積內(nèi)存消耗，從而導(dǎo)致 GPU顯存使用量不斷增加；

??????? 由于博主只使用 VGG16 預(yù)訓(xùn)練模型進(jìn)行分類預(yù)測，不需要訓(xùn)練和反向傳播更新參數(shù)，所以不用構(gòu)建完整的計(jì)算圖。

2--問題解決

??????? 在推理代碼中增加以下指令，表明當(dāng)前計(jì)算不需要進(jìn)行反向傳播，即強(qiáng)制不進(jìn)行完整計(jì)算圖的構(gòu)建：

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

??????? 完整代碼：文章來源地址http://www.zghlxwxcb.cn/news/detail-582128.html

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) # 對預(yù)訓(xùn)練模型進(jìn)行處理
 
        self.linear1 = nn.Linear(in_features = 4096, out_features = self.linear_Class)
 
    def process_backbone(self, model):
 
        # 固定預(yù)訓(xùn)練模型的參數(shù)
        for param in model.parameters():
            param.requires_grad = False
        
        # 刪除最后預(yù)測層    
        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!")

到了這里，關(guān)于深度學(xué)習(xí)筆記--解決GPU顯存使用量不斷增加的問題的文章就介紹完了。如果您還想了解更多內(nèi)容，請?jiān)谟疑辖撬阉鱐OY模板網(wǎng)以前的文章或繼續(xù)瀏覽下面的相關(guān)文章，希望大家以后多多支持TOY模板網(wǎng)！