File size: 3,460 Bytes
2b21abc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 |
import torch
import torch.nn as nn
class SpatioTemporalLSTMCell(nn.Module):
def __init__(self, in_channel, num_hidden, height, width, filter_size, stride, layer_norm):
super(SpatioTemporalLSTMCell, self).__init__()
self.num_hidden = num_hidden
self.padding = filter_size // 2
self._forget_bias = 1.0
if layer_norm:
self.conv_x = nn.Sequential(
nn.Conv2d(in_channel, num_hidden * 7, kernel_size=filter_size,
stride=stride, padding=self.padding, bias=False),
nn.LayerNorm([num_hidden * 7, height, width])
)
self.conv_h = nn.Sequential(
nn.Conv2d(num_hidden, num_hidden * 4, kernel_size=filter_size,
stride=stride, padding=self.padding, bias=False),
nn.LayerNorm([num_hidden * 4, height, width])
)
self.conv_m = nn.Sequential(
nn.Conv2d(num_hidden, num_hidden * 3, kernel_size=filter_size,
stride=stride, padding=self.padding, bias=False),
nn.LayerNorm([num_hidden * 3, height, width])
)
self.conv_o = nn.Sequential(
nn.Conv2d(num_hidden * 2, num_hidden, kernel_size=filter_size,
stride=stride, padding=self.padding, bias=False),
nn.LayerNorm([num_hidden, height, width])
)
else:
self.conv_x = nn.Sequential(
nn.Conv2d(in_channel, num_hidden * 7, kernel_size=filter_size,
stride=stride, padding=self.padding, bias=False),
)
self.conv_h = nn.Sequential(
nn.Conv2d(num_hidden, num_hidden * 4, kernel_size=filter_size,
stride=stride, padding=self.padding, bias=False),
)
self.conv_m = nn.Sequential(
nn.Conv2d(num_hidden, num_hidden * 3, kernel_size=filter_size,
stride=stride, padding=self.padding, bias=False),
)
self.conv_o = nn.Sequential(
nn.Conv2d(num_hidden * 2, num_hidden, kernel_size=filter_size,
stride=stride, padding=self.padding, bias=False),
)
self.conv_last = nn.Conv2d(num_hidden * 2, num_hidden, kernel_size=1,
stride=1, padding=0, bias=False)
def forward(self, x_t, h_t, c_t, m_t):
x_concat = self.conv_x(x_t)
h_concat = self.conv_h(h_t)
m_concat = self.conv_m(m_t)
i_x, f_x, g_x, i_x_prime, f_x_prime, g_x_prime, o_x = \
torch.split(x_concat, self.num_hidden, dim=1)
i_h, f_h, g_h, o_h = torch.split(h_concat, self.num_hidden, dim=1)
i_m, f_m, g_m = torch.split(m_concat, self.num_hidden, dim=1)
i_t = torch.sigmoid(i_x + i_h)
f_t = torch.sigmoid(f_x + f_h + self._forget_bias)
g_t = torch.tanh(g_x + g_h)
c_new = f_t * c_t + i_t * g_t
i_t_prime = torch.sigmoid(i_x_prime + i_m)
f_t_prime = torch.sigmoid(f_x_prime + f_m + self._forget_bias)
g_t_prime = torch.tanh(g_x_prime + g_m)
m_new = f_t_prime * m_t + i_t_prime * g_t_prime
mem = torch.cat((c_new, m_new), 1)
o_t = torch.sigmoid(o_x + o_h + self.conv_o(mem))
h_new = o_t * torch.tanh(self.conv_last(mem))
return h_new, c_new, m_new
|