From 46c2c676ba3a391cc69ce5175833312d51bcae5a Mon Sep 17 00:00:00 2001 From: Stanislaw Adaszewski Date: Fri, 22 May 2020 13:49:09 +0200 Subject: [PATCH] Add type annotations to convolve. --- src/decagon_pytorch/convolve.py | 376 +++++++++++++++++++++++--------- src/decagon_pytorch/layer.py | 2 +- 2 files changed, 273 insertions(+), 105 deletions(-) diff --git a/src/decagon_pytorch/convolve.py b/src/decagon_pytorch/convolve.py index 3b116e7..d50f4b6 100644 --- a/src/decagon_pytorch/convolve.py +++ b/src/decagon_pytorch/convolve.py @@ -1,104 +1,272 @@ -import torch -from .dropout import dropout_sparse, \ - dropout -from .weights import init_glorot - - -class SparseGraphConv(torch.nn.Module): - """Convolution layer for sparse inputs.""" - def __init__(self, in_channels, out_channels, - adjacency_matrix, **kwargs): - super().__init__(**kwargs) - self.in_channels = in_channels - self.out_channels = out_channels - self.weight = init_glorot(in_channels, out_channels) - self.adjacency_matrix = adjacency_matrix - - - def forward(self, x): - x = torch.sparse.mm(x, self.weight) - x = torch.sparse.mm(self.adjacency_matrix, x) - return x - - -class SparseDropoutGraphConvActivation(torch.nn.Module): - def __init__(self, input_dim, output_dim, - adjacency_matrix, keep_prob=1., - activation=torch.nn.functional.relu, - **kwargs): - super().__init__(**kwargs) - self.sparse_graph_conv = SparseGraphConv(input_dim, output_dim, adjacency_matrix) - self.keep_prob = keep_prob - self.activation = activation - - def forward(self, x): - x = dropout_sparse(x, self.keep_prob) - x = self.sparse_graph_conv(x) - x = self.activation(x) - return x - - -class SparseMultiDGCA(torch.nn.Module): - def __init__(self, input_dim, output_dim, - adjacency_matrices, keep_prob=1., - activation=torch.nn.functional.relu, - **kwargs): - super().__init__(**kwargs) - self.output_dim = output_dim - self.sparse_dgca = [ SparseDropoutGraphConvActivation(input_dim, output_dim, adj_mat, keep_prob, activation) for adj_mat in adjacency_matrices ] - - def forward(self, x): - out = torch.zeros(len(x), self.output_dim, dtype=x.dtype) - for f in self.sparse_dgca: - out += f(x) - out = torch.nn.functional.normalize(out, p=2, dim=1) - return out - - -class GraphConv(torch.nn.Module): - def __init__(self, in_channels, out_channels, - adjacency_matrix, **kwargs): - super().__init__(**kwargs) - self.in_channels = in_channels - self.out_channels = out_channels - self.weight = init_glorot(in_channels, out_channels) - self.adjacency_matrix = adjacency_matrix - - def forward(self, x): - x = torch.mm(x, self.weight) - x = torch.mm(self.adjacency_matrix, x) - return x - - -class DropoutGraphConvActivation(torch.nn.Module): - def __init__(self, input_dim, output_dim, - adjacency_matrix, keep_prob=1., - activation=torch.nn.functional.relu, - **kwargs): - super().__init__(**kwargs) - self.graph_conv = GraphConv(input_dim, output_dim, adjacency_matrix) - self.keep_prob = keep_prob - self.activation = activation - - def forward(self, x): - x = dropout(x, keep_prob=self.keep_prob) - x = self.graph_conv(x) - x = self.activation(x) - return x - - -class MultiDGCA(torch.nn.Module): - def __init__(self, input_dim, output_dim, - adjacency_matrices, keep_prob=1., - activation=torch.nn.functional.relu, - **kwargs): - super().__init__(**kwargs) - self.output_dim = output_dim - self.dgca = [ DropoutGraphConvActivation(input_dim, output_dim, adj_mat, keep_prob, activation) for adj_mat in adjacency_matrices ] - - def forward(self, x): - out = torch.zeros(len(x), self.output_dim, dtype=x.dtype) - for f in self.dgca: - out += f(x) - out = torch.nn.functional.normalize(out, p=2, dim=1) - return out +# +# Copyright (C) Stanislaw Adaszewski, 2020 +# License: GPLv3 +# + +""" +This module implements the basic convolutional blocks of Decagon. +Just as a quick reminder, the basic convolution formula here is: + +y = A * (x * W) + +where: + +W is a weight matrix +A is an adjacency matrix +x is a matrix of latent representations of a particular type of neighbors. + +As we have x here twice, a trick is obviously necessary for this to work. +A must be previously normalized with: + +c_{r}^{ij} = 1/sqrt(|N_{r}^{i}| |N_{r}^{j}|) + +or + +c_{r}^{i} = 1/|N_{r}^{i}| + +Let's work through this step by step to convince ourselves that the +formula is correct. + +x = [ + [0, 1, 0, 1], + [1, 1, 1, 0], + [0, 0, 0, 1] +] + +W = [ + [0, 1], + [1, 0], + [0.5, 0.5], + [0.25, 0.75] +] + +A = [ + [0, 1, 0], + [1, 0, 1], + [0, 1, 0] +] + +so the graph looks like this: + +(0) -- (1) -- (2) + +and therefore the representations in the next layer should be: + +h_{0}^{k+1} = c_{r}^{0,1} * h_{1}^{k} * W + c_{r}^{0} * h_{0}^{k} +h_{1}^{k+1} = c_{r}^{0,1} * h_{0}^{k} * W + c_{r}^{2,1} * h_{2}^{k} + + c_{r}^{1} * h_{1}^{k} +h_{2}^{k+1} = c_{r}^{2,1} * h_{1}^{k} * W + c_{r}^{2} * h_{2}^{k} + +In actual Decagon code we can see that that latter part propagating directly +the old representation is gone. I will try to do the same for now. + +So we have to only take care of: + +h_{0}^{k+1} = c_{r}^{0,1} * h_{1}^{k} * W +h_{1}^{k+1} = c_{r}^{0,1} * h_{0}^{k} * W + c_{r}^{2,1} * h_{2}^{k} +h_{2}^{k+1} = c_{r}^{2,1} * h_{1}^{k} * W + +If A is square the Decagon's EdgeMinibatchIterator preprocesses it as follows: + +A = A + eye(len(A)) +rowsum = A.sum(1) +deg_mat_inv_sqrt = diags(power(rowsum, -0.5)) +A = dot(A, deg_mat_inv_sqrt) +A = A.transpose() +A = A.dot(deg_mat_inv_sqrt) + +Let's see what gives in our case: + +A = A + eye(len(A)) + +[ + [1, 1, 0], + [1, 1, 1], + [0, 1, 1] +] + +rowsum = A.sum(1) + +[2, 3, 2] + +deg_mat_inv_sqrt = diags(power(rowsum, -0.5)) + +[ + [1./sqrt(2), 0, 0], + [0, 1./sqrt(3), 0], + [0, 0, 1./sqrt(2)] +] + +A = dot(A, deg_mat_inv_sqrt) + +[ + [ 1/sqrt(2), 1/sqrt(3), 0 ], + [ 1/sqrt(2), 1/sqrt(3), 1/sqrt(2) ], + [ 0, 1/sqrt(3), 1/sqrt(2) ] +] + +A = A.transpose() + +[ + [ 1/sqrt(2), 1/sqrt(2), 0 ], + [ 1/sqrt(3), 1/sqrt(3), 1/sqrt(3) ], + [ 0, 1/sqrt(2), 1/sqrt(2) ] +] + +A = A.dot(deg_mat_inv_sqrt) + +[ + [ 1/sqrt(2) * 1/sqrt(2), 1/sqrt(2) * 1/sqrt(3), 0 ], + [ 1/sqrt(3) * 1/sqrt(2), 1/sqrt(3) * 1/sqrt(3), 1/sqrt(3) * 1/sqrt(2) ], + [ 0, 1/sqrt(2) * 1/sqrt(3), 1/sqrt(2) * 1/sqrt(2) ], +] + +thus: + +[ + [0.5 , 0.40824829, 0. ], + [0.40824829, 0.33333333, 0.40824829], + [0. , 0.40824829, 0.5 ] +] + +This checks out with the 1/sqrt(|N_{r}^{i}| |N_{r}^{j}|) formula. + +Then, we get back to the main calculation: + +y = x * W +y = A * y + +y = x * W + +[ + [ 1.25, 0.75 ], + [ 1.5 , 1.5 ], + [ 0.25, 0.75 ] +] + +y = A * y + +[ + 0.5 * [ 1.25, 0.75 ] + 0.40824829 * [ 1.5, 1.5 ], + 0.40824829 * [ 1.25, 0.75 ] + 0.33333333 * [ 1.5, 1.5 ] + 0.40824829 * [ 0.25, 0.75 ], + 0.40824829 * [ 1.5, 1.5 ] + 0.5 * [ 0.25, 0.75 ] +] + +that is: + +[ + [1.23737243, 0.98737244], + [1.11237243, 1.11237243], + [0.73737244, 0.98737244] +]. + +All checks out nicely, good. + +""" + + +import torch +from .dropout import dropout_sparse, \ + dropout +from .weights import init_glorot +from typing import List, Callable + + +class SparseGraphConv(torch.nn.Module): + """Convolution layer for sparse inputs.""" + def __init__(self, in_channels: int, out_channels: int, + adjacency_matrix: torch.Tensor, **kwargs) -> None: + super().__init__(**kwargs) + self.in_channels = in_channels + self.out_channels = out_channels + self.weight = init_glorot(in_channels, out_channels) + self.adjacency_matrix = adjacency_matrix + + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = torch.sparse.mm(x, self.weight) + x = torch.sparse.mm(self.adjacency_matrix, x) + return x + + +class SparseDropoutGraphConvActivation(torch.nn.Module): + def __init__(self, input_dim: int, output_dim: int, + adjacency_matrix: torch.Tensor, keep_prob: float=1., + activation: Callable[[torch.Tensor], torch.Tensor]=torch.nn.functional.relu, + **kwargs) -> None: + super().__init__(**kwargs) + self.sparse_graph_conv = SparseGraphConv(input_dim, output_dim, adjacency_matrix) + self.keep_prob = keep_prob + self.activation = activation + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = dropout_sparse(x, self.keep_prob) + x = self.sparse_graph_conv(x) + x = self.activation(x) + return x + + +class SparseMultiDGCA(torch.nn.Module): + def __init__(self, input_dim: List[int], output_dim: int, + adjacency_matrices: List[torch.Tensor], keep_prob: float=1., + activation: Callable[[torch.Tensor], torch.Tensor]=torch.nn.functional.relu, + **kwargs) -> None: + super().__init__(**kwargs) + self.output_dim = output_dim + self.sparse_dgca = [ SparseDropoutGraphConvActivation(input_dim, output_dim, adj_mat, keep_prob, activation) for adj_mat in adjacency_matrices ] + + def forward(self, x: List[torch.Tensor]) -> List[torch.Tensor]: + out = torch.zeros(len(x), self.output_dim, dtype=x.dtype) + for f in self.sparse_dgca: + out += f(x) + out = torch.nn.functional.normalize(out, p=2, dim=1) + return out + + +class GraphConv(torch.nn.Module): + def __init__(self, in_channels: int, out_channels: int, + adjacency_matrix: torch.Tensor, **kwargs) -> None: + super().__init__(**kwargs) + self.in_channels = in_channels + self.out_channels = out_channels + self.weight = init_glorot(in_channels, out_channels) + self.adjacency_matrix = adjacency_matrix + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = torch.mm(x, self.weight) + x = torch.mm(self.adjacency_matrix, x) + return x + + +class DropoutGraphConvActivation(torch.nn.Module): + def __init__(self, input_dim: int, output_dim: int, + adjacency_matrix: torch.Tensor, keep_prob: float=1., + activation: Callable[[torch.Tensor], torch.Tensor]=torch.nn.functional.relu, + **kwargs) -> None: + super().__init__(**kwargs) + self.graph_conv = GraphConv(input_dim, output_dim, adjacency_matrix) + self.keep_prob = keep_prob + self.activation = activation + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = dropout(x, keep_prob=self.keep_prob) + x = self.graph_conv(x) + x = self.activation(x) + return x + + +class MultiDGCA(torch.nn.Module): + def __init__(self, input_dim: List[int], output_dim: int, + adjacency_matrices: List[torch.Tensor], keep_prob: float=1., + activation: Callable[[torch.Tensor], torch.Tensor]=torch.nn.functional.relu, + **kwargs) -> None: + super().__init__(**kwargs) + self.output_dim = output_dim + self.dgca = [ DropoutGraphConvActivation(input_dim, output_dim, adj_mat, keep_prob, activation) for adj_mat in adjacency_matrices ] + + def forward(self, x: List[torch.Tensor]) -> List[torch.Tensor]: + out = torch.zeros(len(x), self.output_dim, dtype=x.dtype) + for f in self.dgca: + out += f(x) + out = torch.nn.functional.normalize(out, p=2, dim=1) + return out diff --git a/src/decagon_pytorch/layer.py b/src/decagon_pytorch/layer.py index 97c7c2a..a75c005 100644 --- a/src/decagon_pytorch/layer.py +++ b/src/decagon_pytorch/layer.py @@ -21,7 +21,7 @@ import torch -from .convole import SparseMultiDGCA +from .convolve import SparseMultiDGCA class InputLayer(torch.nn.Module):