Homec4science
Diffusion project14 (master) Browse

R6590/lstm.py
221b9af50125master

Last Change
Raw File

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, dropout = 0.2)
# 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, future = 0):
'''
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
'''
outputs = []
hidden_states = Variable(torch.zeros(1,1,self.hidden_size))
cell_states = Variable(torch.zeros(1,1,self.hidden_size))
#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)
outputs += [output]
for i in range(future): #predict future
new_out = self.relu(self.input_embedding_layer(output))
h_peds, c_peds = self.lstm(new_out, (hidden_states, cell_states))
output = self.output_layer(h_peds)
outputs += [output]
outputs = torch.cat(outputs, 0)
return outputs
c4science · Help