|
|
import torch |
|
|
import torch.nn as nn |
|
|
import math |
|
|
|
|
|
|
|
|
class MAUCell(nn.Module): |
|
|
|
|
|
def __init__(self, in_channel, num_hidden, height, width, filter_size, stride, tau, cell_mode): |
|
|
super(MAUCell, self).__init__() |
|
|
|
|
|
self.num_hidden = num_hidden |
|
|
|
|
|
self.padding = filter_size // 2 |
|
|
self.cell_mode = cell_mode |
|
|
self.d = num_hidden * height * width |
|
|
self.tau = tau |
|
|
self.states = ['residual', 'normal'] |
|
|
if not self.cell_mode in self.states: |
|
|
raise AssertionError |
|
|
self.conv_t = nn.Sequential( |
|
|
nn.Conv2d(in_channel, 3 * num_hidden, kernel_size=filter_size, |
|
|
stride=stride, padding=self.padding), |
|
|
nn.LayerNorm([3 * num_hidden, height, width]) |
|
|
) |
|
|
self.conv_t_next = nn.Sequential( |
|
|
nn.Conv2d(in_channel, num_hidden, kernel_size=filter_size, |
|
|
stride=stride, padding=self.padding), |
|
|
nn.LayerNorm([num_hidden, height, width]) |
|
|
) |
|
|
self.conv_s = nn.Sequential( |
|
|
nn.Conv2d(num_hidden, 3 * num_hidden, kernel_size=filter_size, |
|
|
stride=stride, padding=self.padding), |
|
|
nn.LayerNorm([3 * num_hidden, height, width]) |
|
|
) |
|
|
self.conv_s_next = nn.Sequential( |
|
|
nn.Conv2d(num_hidden, num_hidden, kernel_size=filter_size, |
|
|
stride=stride, padding=self.padding), |
|
|
nn.LayerNorm([num_hidden, height, width]) |
|
|
) |
|
|
self.softmax = nn.Softmax(dim=0) |
|
|
|
|
|
def forward(self, T_t, S_t, t_att, s_att): |
|
|
s_next = self.conv_s_next(S_t) |
|
|
t_next = self.conv_t_next(T_t) |
|
|
weights_list = [] |
|
|
for i in range(self.tau): |
|
|
weights_list.append((s_att[i] * s_next).sum(dim=(1, 2, 3)) / math.sqrt(self.d)) |
|
|
weights_list = torch.stack(weights_list, dim=0) |
|
|
weights_list = torch.reshape(weights_list, (*weights_list.shape, 1, 1, 1)) |
|
|
weights_list = self.softmax(weights_list) |
|
|
T_trend = t_att * weights_list |
|
|
T_trend = T_trend.sum(dim=0) |
|
|
t_att_gate = torch.sigmoid(t_next) |
|
|
T_fusion = T_t * t_att_gate + (1 - t_att_gate) * T_trend |
|
|
T_concat = self.conv_t(T_fusion) |
|
|
S_concat = self.conv_s(S_t) |
|
|
t_g, t_t, t_s = torch.split(T_concat, self.num_hidden, dim=1) |
|
|
s_g, s_t, s_s = torch.split(S_concat, self.num_hidden, dim=1) |
|
|
T_gate = torch.sigmoid(t_g) |
|
|
S_gate = torch.sigmoid(s_g) |
|
|
T_new = T_gate * t_t + (1 - T_gate) * s_t |
|
|
S_new = S_gate * s_s + (1 - S_gate) * t_s |
|
|
|
|
|
if self.cell_mode == 'residual': |
|
|
S_new = S_new + S_t |
|
|
return T_new, S_new |
|
|
|
