用于从 windows 到 Linux 传

1. try_hook.py

import numpy as np

import torch
import torchvision
from PIL import Image
from torchvision import transforms as T

import matplotlib.pyplot as plt

feature_extractor = torchvision.models.resnet34(pretrained=True)
if torch.cuda.is_available():
    feature_extractor.cuda()
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')

class SaveOutput:
    def __init__(self):
        self.outputs = []

    def __call__(self, module, module_in, module_out):
        self.outputs.append(module_out)

    def clear(self):
        self.outputs = []

save_output = SaveOutput()
hook_handles = []
for layer in feature_extractor.modules():
    if isinstance(layer, torch.nn.Conv2d):
        handle = layer.register_forward_hook(save_output)
        hook_handles.append(handle)
image = Image.open('cat.jpg')
transform = T.Compose([T.Resize((224, 224)), T.ToTensor()])
X = transform(image).unsqueeze(dim=0).to(device)
out = feature_extractor(X)
print(len(save_output.outputs))
a_list = [0, 1, 6, 15, 28, 35]
for i in a_list:
    print(save_output.outputs[i].cpu().detach().squeeze(0).shape)

def grid_gray_image(imgs, each_row: int):  # each_row 8
    '''
    imgs shape: batch * size (e.g., 64x32x32, 64 is the number of the gray images, and (32, 32) is the size of each gray image)
    '''
    row_num = imgs.shape[0]//each_row  # imgs.shape[0] 是通道数
    for i in range(row_num):
        img = imgs[i*each_row]
        img = (img - img.min()) / (img.max() - img.min())  # 归一化
        for j in range(1, each_row):
            tmp_img = imgs[i*each_row+j]
            tmp_img = (tmp_img - tmp_img.min()) / (tmp_img.max() - tmp_img.min())  # 归一化
            img = np.hstack((img, tmp_img))
        if i == 0:
            ans = img
        else:
            ans = np.vstack((ans, img))
    return ans
img0 = save_output.outputs[0].cpu().detach().squeeze(0)  # tensor:(64,112,112)
img0 = grid_gray_image(img0.numpy(), 8)  # ndarray:(896,896)
img1 = save_output.outputs[1].cpu().detach().squeeze(0)
img1 = grid_gray_image(img1.numpy(), 8)
img6 = save_output.outputs[6].cpu().detach().squeeze(0)
img6 = grid_gray_image(img6.numpy(), 8)
img15 = save_output.outputs[15].cpu().detach().squeeze(0)
img15 = grid_gray_image(img15.numpy(), 16)
img29 = save_output.outputs[28].cpu().detach().squeeze(0)
img29 = grid_gray_image(img29.numpy(), 16)
# plt.figure(figsize=(15, 15))
# plt.imshow(img0, cmap='gray')
# plt.savefig('./result/0.jpg')
# plt.figure(figsize=(15, 15))
# plt.imshow(img1, cmap='gray')
# plt.savefig('./result/1.jpg')
# plt.show()

2. try_colormap.py

import cv2

def cv_show(name, img):
    cv2.imshow('name', img)
    cv2.waitKey(30000)
    cv2.destroyAllWindows()


img1 = cv2.imread('result/0.jpg', cv2.IMREAD_GRAYSCALE)
print(img1.shape)
cv_show(1, img1)
img_color = cv2.applyColorMap(img1, cv2.COLORMAP_JET)
cv_show(2, img_color)