import torchfrom torchvision.models import resnet18import torch.nn as nnfrom torchvision import transforms
import matplotlib.pyplot as plt

def viz(module, input):    x = input[0][0]    #最多顯示4張圖    min_num = np.minimum(4, x.size()[0])    for i in range(min_num):        plt.subplot(1, 4, i+1)        plt.imshow(x[i])    plt.show()

import cv2import numpy as npdef main():    t = transforms.Compose([transforms.ToPILImage(),                            transforms.Resize((224, 224)),                            transforms.ToTensor(),                            transforms.Normalize(mean=[0.485, 0.456, 0.406],                                                 std=[0.229, 0.224, 0.225])                            ])
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    model = resnet18(pretrained=True).to(device)    for name, m in model.named_modules():        # if not isinstance(m, torch.nn.ModuleList) and        #         not isinstance(m, torch.nn.Sequential) and        #         type(m) in torch.nn.__dict__.values():        # 這里只對卷積層的feature map進行顯示        if isinstance(m, torch.nn.Conv2d):            m.register_forward_pre_hook(viz)    img = cv2.imread('/Users/edgar/Desktop/cat.jpeg')    img = t(img).unsqueeze(0).to(device)    with torch.no_grad():        model(img)
if __name__ == '__main__':    main()

 inter_feature = {} inter_gradient = {} def make_hook(name, flag):     if flag == 'forward':         def hook(m, input, output):             inter_feature[name] = input         return hook     elif flag == 'backward':         def hook(m, input, output):             inter_gradient[name] = output         return hook     else:         assert Falsem.register_forward_hook(make_hook(name, 'forward'))m.register_backward_hook(make_hook(name, 'backward'))

output = model(input)  # achieve intermediate featureloss = criterion(output, target)loss.backward()  # achieve backward intermediate gradients