I am trying to visualize my LSTM model, which looks like this:
class SequenceModel(nn.Module):
def __init__(self, n_features, n_classes, n_hidden=100, n_layers=3):
super().__init__()
self.lstm = nn.LSTM(
input_size = n_features,
hidden_size = n_hidden,
num_layers = n_layers,
batch_first = True,
dropout = 0.945
)
self.classifier = nn.Linear(n_hidden, n_classes)
def forward(self,x):
self.lstm.flatten_parameters()
_, (hidden,_) = self.lstm(x)
out = hidden [-1]
return self.classifier(out)
(model): SequenceModel(
(lstm): LSTM(3, 100, num_layers=3, batch_first=True, dropout=0.945)
(classifier): Linear(in_features=100, out_features=2, bias=True)
)
(criterion): CrossEntropyLoss()
)
I know usually you do this like this:
To write the computational graph we will be using add_graph() method. add_graph requires two arguments
- The model
- A sample image for the same shape as that of the input to track how it changes as it passes through the network
So i tried it with this code:
from torch.utils.tensorboard import SummaryWriter
sample=torch.rand((3,2,3))
writer = SummaryWriter()
writer.add_graph(model, sample)
writer.close()
But I am getting this error:
Only tensors, lists, tuples of tensors, or dictionary of tensors can be output from traced functions
Error occurs, No graph saved
I can´t figure out what the problem is, has anyone an idea? Thanks in advance.