深度学习笔记

您所在的位置：网站首页 › pytorchgpu要求 › 深度学习笔记

深度学习笔记

#深度学习笔记| 来源: 网络整理| 查看: 265

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

【本文地址】

深度学习笔记

深度学习笔记

今日新闻

推荐新闻