pytorch保存一张图的所有特征图

刀么克瑟拉莫

已于 2022-08-30 11:55:04 修改

阅读量959

点赞数

分类专栏： pytorch 文章标签： pytorch 深度学习 python

于 2022-08-30 11:54:53 首次发布

本文链接：https://blog.csdn.net/random_repick/article/details/126600962

版权

pytorch 专栏收录该内容

17 篇文章 0 订阅

订阅专栏

需要先了解所用网络的结构

import torch.nn as nn
import torch

class Simple(nn.Module):
	def __init__(self):
		super().__init__()
		self.feature = nn.Sequential(
			nn.Conv2d(3, 16, 3, 1, 1, bias=True),
			nn.MaxPool2d(2),
			nn.Conv2d(16, 32, 3, 1, 1, bias=True),
			nn.MaxPool2d(2),
			nn.Conv2d(32, 64, 3, 1, 1, bias=True),
			nn.MaxPool2d(2),
			nn.Conv2d(64, 128, 3, 1, 1, bias=True),
			nn.MaxPool2d(2),
			nn.Conv2d(128, 256, 3, 1, 1, bias=True),
			nn.MaxPool2d(2),
			nn.Conv2d(256, 512, 3, 1, 1, bias=True)
		)
		self.fc = nn.Sequential(
			nn.Linear(512*16, 1000),
			nn.ReLU(),
			nn.Linear(1000, 100),
			nn.ReLU(),
			nn.Linear(100, 6)
		)
		
	def forward(self, x):
		f = self.feature(x)
		f = f.view(-1, 512*16)
		return self.fc(f)

加载模型，预处理图片，记得要保证存储目录存在哦～

net=Simple()
paras=torch.load('', map_location='cuda')
net.load_state_dict(paras)
net.cuda().eval()
raw = cv2.imread('')
im = cv2.resize(raw, (150, 150))
im = im / 255.0
im = np.transpose(im, (2,0,1))
im = im.reshape(1,3,150,150)
im = torch.from_numpy(im).float().cuda()
for ch in net.feature.children():
		im = ch(im)
		b, c, h, w = im.shape
		for ci in range(c):
			fm = im[0][ci]
			ma = torch.max(fm)
			mi = torch.min(fm)
			fm = 255 * (fm - mi) / (ma - mi)
			fm = fm.cpu().detach().numpy().astype('uint8')
			fm = fm.reshape(h, w, 1)
			cv2.imwrite('log/feature/{}-{}.jpg'.format(h, ci), fm)