Я пытаюсь записать свою модель в tensorboard
с помощью следующего кода:
model = SimpleLSTM(4, HIDDEN_DIM, HIDDEN_LAYERS, 1, BATCH_SIZE, device)
writer = tb.SummaryWriter(log_dir=tb_path)
sample_data = iter(trainloader).next()[0]
writer.add_graph(model, sample_data.to(device))
Я получаю сообщение об ошибке: TypeError: forward() missing 1 required positional argument: 'batch_size'
Моя модель выглядит так:
class SimpleLSTM(nn.Module):
def __init__(self, input_dims, hidden_units, hidden_layers, out, batch_size, device):
super(SimpleLSTM, self).__init__()
self.input_dims = input_dims
self.hidden_units = hidden_units
self.hidden_layers = hidden_layers
self.batch_size = batch_size
self.device = device
self.lstm = nn.LSTM(self.input_dims, self.hidden_units, self.hidden_layers,
batch_first=True, bidirectional=False)
self.output_layer = nn.Linear(self.hidden_units, out)
def init_hidden(self, batch_size):
hidden = torch.rand(self.hidden_layers, batch_size, self.hidden_units, device=self.device, dtype=torch.float32)
cell = torch.rand(self.hidden_layers, batch_size, self.hidden_units, device=self.device, dtype=torch.float32)
hidden = nn.init.xavier_normal_(hidden)
cell = nn.init.xavier_normal_(cell)
return (hidden, cell)
def forward(self, input, batch_size):
hidden = self.init_hidden(batch_size) incomplete batch
lstm_out, (h_n, c_n) = self.lstm(input, hidden)
raw_out = self.output_layer(h_n[-1])
return raw_out
Как записать эту модель на TensorBoard?