Page MenuHomec4science
Files F66217682

lstm.py
No OneTemporary
Actions

File Metadata

Created
Sun, Jun 9, 02:06

lstm.py
View Options

import torch
import torch.nn as nn
from torch.autograd import Variable
class LSTMmodel(nn.Module):
'''
Class representing the LSTM model
'''
def __init__(self):
super(LSTMmodel,self).__init__()
# Store required sizes
self.hidden_size = 128
#self.grid_size = args.grid_size
self.embedding_size = 64
#self.pooling_size = args.pooling_size #pooling window
self.input_size = 2
self.output_size = 2 #parameters of bivariate distribution
#self.neighborhood_size = args.neighborhood_size
# The LSTM cell. (Social LSTM) embedding size = 64
self.lstm= nn.LSTM(self.embedding_size, self.hidden_size)
# Linear layer to embed the input position into LSTM
self.input_embedding_layer = nn.Linear(self.input_size, self.embedding_size)
# Linear layer to embed the social tensor
#self.tensor_embedding_layer = nn.Linear(self.neighborhood_size*self.neighborhood_size*self.hidden_size, self.embedding_size)
# Linear layer to map the hidden state of LSTM to output
self.output_layer = nn.Linear(self.hidden_size, self.output_size)
# ReLU and dropout unit
self.relu = nn.ReLU()
#self.dropout = nn.Dropout(0.5)
def forward(self, peds, hidden_states, cell_states):
'''
Forward pass for the model
params:
peds: coordinates of pedestrians in frame
hidden_states: Hidden states of the pedestrians
cell_states: Cell states of the peds
returns:
weigths: Outputs corresponding to bivariate Gaussian distributions
hidden_states
cell_states
'''
#input = self.dropout(self.relu(self.input_embedding_layer(peds)))
input = self.relu(self.input_embedding_layer(peds))
h_peds, c_peds = self.lstm(input, (hidden_states, cell_states))
output = self.output_layer(h_peds)
#weigths = torch.cat((weights,output),0)
return output, h_peds, c_peds

Event Timeline

c4science · Help