79 lines
No EOL
3 KiB
Python
79 lines
No EOL
3 KiB
Python
import torch.nn as nn
|
|
import torch
|
|
from torch.nn.parameter import Parameter
|
|
import math
|
|
|
|
|
|
class graph_convolution(nn.Module):
|
|
def __init__(self, in_features, out_features, node_n = 21, seq_len = 40, bias=True):
|
|
super().__init__()
|
|
|
|
self.temporal_graph_weights = Parameter(torch.FloatTensor(seq_len, seq_len))
|
|
self.feature_weights = Parameter(torch.FloatTensor(in_features, out_features))
|
|
self.spatial_graph_weights = Parameter(torch.FloatTensor(node_n, node_n))
|
|
|
|
if bias:
|
|
self.bias = Parameter(torch.FloatTensor(seq_len))
|
|
|
|
self.reset_parameters()
|
|
|
|
def reset_parameters(self):
|
|
stdv = 1. / math.sqrt(self.spatial_graph_weights.size(1))
|
|
self.feature_weights.data.uniform_(-stdv, stdv)
|
|
self.temporal_graph_weights.data.uniform_(-stdv, stdv)
|
|
self.spatial_graph_weights.data.uniform_(-stdv, stdv)
|
|
if self.bias is not None:
|
|
self.bias.data.uniform_(-stdv, stdv)
|
|
|
|
def forward(self, input):
|
|
y = torch.matmul(input, self.temporal_graph_weights)
|
|
y = torch.matmul(y.permute(0, 3, 2, 1), self.feature_weights)
|
|
y = torch.matmul(self.spatial_graph_weights, y).permute(0, 3, 2, 1).contiguous()
|
|
|
|
if self.bias is not None:
|
|
return (y + self.bias)
|
|
else:
|
|
return y
|
|
|
|
|
|
class residual_graph_convolution(nn.Module):
|
|
def __init__(self, features, node_n=21, seq_len = 40, bias=True, p_dropout=0.3):
|
|
super().__init__()
|
|
|
|
self.gcn = graph_convolution(features, features, node_n=node_n, seq_len=seq_len, bias=bias)
|
|
self.ln = nn.LayerNorm([features, node_n, seq_len], elementwise_affine=True)
|
|
self.act_f = nn.Tanh()
|
|
self.dropout = nn.Dropout(p_dropout)
|
|
|
|
def forward(self, x):
|
|
|
|
y = self.gcn(x)
|
|
y = self.ln(y)
|
|
y = self.act_f(y)
|
|
y = self.dropout(y)
|
|
|
|
return y + x
|
|
|
|
|
|
class graph_convolution_network(nn.Module):
|
|
def __init__(self, in_features, latent_features, node_n=21, seq_len=40, p_dropout=0.3, residual_gcns_num=1):
|
|
super().__init__()
|
|
self.residual_gcns_num = residual_gcns_num
|
|
self.seq_len = seq_len
|
|
|
|
self.start_gcn = graph_convolution(in_features=in_features, out_features=latent_features, node_n=node_n, seq_len=seq_len)
|
|
|
|
self.residual_gcns = []
|
|
for i in range(residual_gcns_num):
|
|
self.residual_gcns.append(residual_graph_convolution(features=latent_features, node_n=node_n, seq_len=seq_len*2, p_dropout=p_dropout))
|
|
self.residual_gcns = nn.ModuleList(self.residual_gcns)
|
|
|
|
def forward(self, x):
|
|
y = self.start_gcn(x)
|
|
|
|
y = torch.cat((y, y), dim=3)
|
|
for i in range(self.residual_gcns_num):
|
|
y = self.residual_gcns[i](y)
|
|
y = y[:, :, :, :self.seq_len]
|
|
|
|
return y |