I’m not quite sure what your goal is, but you could look into hooks. Below is a complete example that prints the shapes of the input and output tensors for each module.
That being said, it’s not only the shape that matters. Particularly when using view() it’s easy to get the right shape but to mess up the tensors; cf. this post.
import torch
import torch.nn as nn
class Hook():
def __init__(self, module, backward=False):
self.module = module
if backward==False:
self.hook = module.register_forward_hook(self.hook_fn)
else:
self.hook = module.register_backward_hook(self.hook_fn)
def hook_fn(self, module, input, output):
self.input = input
self.output = output
def close(self):
self.hook.remove()
class SimpleLSTM(nn.Module):
def __init__(self, vocab_size, embed_size, hidden_size, output_size):
super().__init__()
self.embed = nn.Embedding(vocab_size, embed_size)
self.lstm = nn.LSTM(embed_size, hidden_size, batch_first=False)
self.net = nn.Sequential(
nn.Linear(hidden_size, 64),
nn.ReLU(),
nn.Linear(64, output_size),
nn.LogSoftmax(dim=1)
)
def forward(self, X):
out = self.embed(X)
out = out.transpose(0, 1)
out, (h, c) = self.lstm(out)
log_probs = self.net(out[-1])
return log_probs
# Instantiate the model (simple LSTM classifier)
model = SimpleLSTM(100, 300, 512, 3)
# Attache a forward hook to each layer
forward_hooks = [Hook(layer[1]) for layer in list(model.named_modules())]
# Create random input patch and pump it through the model
x = torch.randint(vocab_size, (batch_size, seq_len))
out = model(x)
# Print input and output shapes of each layer
for hook in forward_hooks:
if isinstance(hook.input, torch.Tensor):
input_shape = hook.input.shape
else:
input_shape = hook.input[0].shape
if isinstance(hook.output, torch.Tensor):
output_shape = hook.output.shape
else:
output_shape = hook.output[0].shape
print("Module:", hook.module)
print("Input shape:\t", input_shape)
print("Output shape:\t", output_shape)
print('---'*17)