qq7835144的博客

Pytorch多GPU训练时使用hook提取模型中间层输出时与模型输入张量不在同一个GPU上的解决办法 通常对于单卡训练的模型，使用hook可以较为方便地提取出模型中间层输出。 例如我们想要获取自定义模型DBL中的conv2d的输出，可以先打印出这个网络，获取到conv2d在模型中的次序，然后使用for循环确定其位置并注册hook。 参考https://www.jianshu.com/p/0a270d63aca9 import torch import torch.nn as nn class CBL(n