目录
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!")
