I created my own sequential that has many modules in it. I was able to create a module PrintLayer and I placed it after the output i want to save.
class PrintLayer(nn.Module):
def init(self):
super(PrintLayer, self).init()
def forward(self, x):
print(x.shape)
return x
It could print the shape and value of the output of a specific layer i want to observe during the forward pass, but how exactly do I save that value so I could retrieve it after the forward pass. Thanks!
Would it be feasible to set the hook in the __init__ of some sub-module, e.g. UnetSkipConnectionBlock?
Otherwise, you would need to address your layers from top to bottom.
Yeah I can add stuff to the submodule, but I’m unsure how to add the hook there since each submodule is a Sequential and isn’t created like the example you have.
Yes, then I would use the method I’ve linked in the other post.
Note that you might want to remove the .detach in hook, if you want to backpropagate.