| 94 | |
| 95 | |
| 96 | class MTGODE(nn.Module): |
| 97 | |
| 98 | def __init__(self, buildA_true, num_nodes, device, predefined_A=None, static_feat=None, dropout=0.3, |
| 99 | subgraph_size=20, node_dim=40, dilation_exponential=1, conv_channels=32, end_channels=128, |
| 100 | seq_length=12, in_dim=2, out_dim=12, tanhalpha=3, method_1='euler', time_1=1.2, step_size_1=0.4, |
| 101 | method_2='euler', time_2=1.0, step_size_2=0.25, alpha=1.0, rtol=1e-4, atol=1e-3, adjoint=False, |
| 102 | perturb=False, ln_affine=True): |
| 103 | |
| 104 | super(MTGODE, self).__init__() |
| 105 | |
| 106 | if method_1 == 'euler': |
| 107 | self.integration_time = time_1 |
| 108 | self.estimated_nfe = round(self.integration_time / step_size_1) |
| 109 | elif method_1 == 'rk4': |
| 110 | self.integration_time = time_1 |
| 111 | self.estimated_nfe = round(self.integration_time / (step_size_1 / 4.0)) |
| 112 | else: |
| 113 | raise ValueError("Oops! Temporal ODE solver is invaild.") |
| 114 | |
| 115 | self.buildA_true = buildA_true |
| 116 | self.num_nodes = num_nodes |
| 117 | self.dropout = dropout |
| 118 | self.predefined_A = predefined_A |
| 119 | self.seq_length = seq_length |
| 120 | self.ln_affine = ln_affine |
| 121 | self.adjoint = adjoint |
| 122 | |
| 123 | self.start_conv = nn.Conv2d(in_channels=in_dim, out_channels=conv_channels, kernel_size=(1, 1)) |
| 124 | |
| 125 | self.gc = graph_constructor(num_nodes, subgraph_size, node_dim, device, alpha=tanhalpha, static_feat=static_feat) |
| 126 | self.idx = torch.arange(self.num_nodes).to(device) |
| 127 | |
| 128 | max_kernel_size = 7 |
| 129 | if dilation_exponential > 1: |
| 130 | self.receptive_field = int(1 + (max_kernel_size - 1) * (dilation_exponential**self.estimated_nfe - 1) / (dilation_exponential - 1)) |
| 131 | else: |
| 132 | self.receptive_field = self.estimated_nfe * (max_kernel_size - 1) + 1 |
| 133 | |
| 134 | if ln_affine: |
| 135 | self.affine_weight = nn.Parameter(torch.Tensor(*(conv_channels, self.num_nodes))) # C*H |
| 136 | self.affine_bias = nn.Parameter(torch.Tensor(*(conv_channels, self.num_nodes))) # C*H |
| 137 | |
| 138 | self.ODE = ODEBlock(ODEFunc(STBlock(receptive_field=self.receptive_field, dilation=dilation_exponential, |
| 139 | hidden_channels=conv_channels, dropout=self.dropout, method=method_2, |
| 140 | time=time_2, step_size=step_size_2, alpha=alpha, rtol=rtol, atol=atol, |
| 141 | adjoint=False, perturb=perturb)), |
| 142 | method_1, step_size_1, rtol, atol, adjoint, perturb) |
| 143 | |
| 144 | self.end_conv_0 = nn.Conv2d(in_channels=conv_channels, out_channels=end_channels//2, kernel_size=(1, 1), bias=True) |
| 145 | self.end_conv_1 = nn.Conv2d(in_channels=end_channels//2, out_channels=end_channels, kernel_size=(1, 1), bias=True) |
| 146 | self.end_conv_2 = nn.Conv2d(in_channels=end_channels, out_channels=out_dim, kernel_size=(1, 1), bias=True) |
| 147 | |
| 148 | if ln_affine: |
| 149 | self.reset_parameters() |
| 150 | |
| 151 | def reset_parameters(self): |
| 152 | init.ones_(self.affine_weight) |
| 153 | init.zeros_(self.affine_bias) |