r"""An implementation of the Diffusion Convolution Layer. For details see: `"Diffusion Convolutional Recurrent Neural Network: Data-Driven Traffic Forecasting" `_ Args: in_channels (int): Number of input features. out_channels (int): Num
| 299 | |
| 300 | |
| 301 | class DConv(MessagePassing): |
| 302 | r"""An implementation of the Diffusion Convolution Layer. |
| 303 | For details see: `"Diffusion Convolutional Recurrent Neural Network: |
| 304 | Data-Driven Traffic Forecasting" <https://arxiv.org/abs/1707.01926>`_ |
| 305 | |
| 306 | Args: |
| 307 | in_channels (int): Number of input features. |
| 308 | out_channels (int): Number of output features. |
| 309 | K (int): Filter size :math:`K`. |
| 310 | bias (bool, optional): If set to :obj:`False`, the layer |
| 311 | will not learn an additive bias (default :obj:`True`). |
| 312 | |
| 313 | """ |
| 314 | |
| 315 | def __init__(self, in_channels, out_channels, K, bias=True): |
| 316 | super(DConv, self).__init__(aggr="add", flow="source_to_target") |
| 317 | assert K > 0 |
| 318 | self.in_channels = in_channels |
| 319 | self.out_channels = out_channels |
| 320 | self.weight = torch.nn.Parameter(torch.Tensor(2, K, in_channels, out_channels)) |
| 321 | |
| 322 | if bias: |
| 323 | self.bias = torch.nn.Parameter(torch.Tensor(out_channels)) |
| 324 | else: |
| 325 | self.register_parameter("bias", None) |
| 326 | |
| 327 | self.__reset_parameters() |
| 328 | |
| 329 | def __reset_parameters(self): |
| 330 | torch.nn.init.xavier_uniform_(self.weight) |
| 331 | if self.bias is not None: torch.nn.init.zeros_(self.bias) |
| 332 | |
| 333 | def message(self, x_j, norm): |
| 334 | return norm.view(-1, 1) * x_j |
| 335 | |
| 336 | def forward( |
| 337 | self, |
| 338 | X: torch.FloatTensor, |
| 339 | edge_index: torch.LongTensor, |
| 340 | edge_weight: torch.FloatTensor, |
| 341 | ) -> torch.FloatTensor: |
| 342 | r"""Making a forward pass. If edge weights are not present the forward pass |
| 343 | defaults to an unweighted graph. |
| 344 | |
| 345 | Arg types: |
| 346 | * **X** (PyTorch Float Tensor) - Node features. |
| 347 | * **edge_index** (PyTorch Long Tensor) - Graph edge indices. |
| 348 | * **edge_weight** (PyTorch Long Tensor, optional) - Edge weight vector. |
| 349 | |
| 350 | Return types: |
| 351 | * **H** (PyTorch Float Tensor) - Hidden state matrix for all nodes. |
| 352 | """ |
| 353 | adj_mat = to_dense_adj(edge_index, edge_attr=edge_weight) |
| 354 | adj_mat = adj_mat.reshape(adj_mat.size(1), adj_mat.size(2)) |
| 355 | deg_out = torch.matmul( |
| 356 | adj_mat, torch.ones(size=(adj_mat.size(0), 1)).to(X.device) |
| 357 | ) |
| 358 | deg_out = deg_out.flatten() |