当需要拿出model的一些中间输出并做bp等操作,可以利用hook机制
对某个module调用register_forward_hook(get)时,会在forward时,自动调用定义的函数get(也可以是其他名字)但参数是model input output。使用input和model,在hook标注的这个module得到的结果会返回到output,然后可以对它进行一些操作
我这里是把一些feature放到了model的一个self.features参数里
def get_output(self):
# x = imgs
self.features = []
def get(model, input