def __init__(self, in_channels, out_channels, dim, kernel_size, is_open_spline=True, degree=1, aggr='mean', root_weight=True, bias=True, **kwargs): super(SplineConv, self).__init__(aggr=aggr, **kwargs) if spline_basis is None: raise ImportError('`SplineConv` requires `torch-spline-conv`.') self.in_channels = in_channels self.out_channels = out_channels self.dim = dim self.degree = degree kernel_size = torch.tensor(repeat(kernel_size, dim), dtype=torch.long) self.register_buffer('kernel_size', kernel_size) is_open_spline = repeat(is_open_spline, dim) is_open_spline = torch.tensor(is_open_spline, dtype=torch.uint8) self.register_buffer('is_open_spline', is_open_spline) K = kernel_size.prod().item() self.weight = Parameter(torch.Tensor(K, in_channels, out_channels)) if root_weight: self.root = Parameter(torch.Tensor(in_channels, out_channels)) else: self.register_parameter('root', None) if bias: self.bias = Parameter(torch.Tensor(out_channels)) else: self.register_parameter('bias', None) self.reset_parameters()
def voxel_grid( pos: Tensor, size: Union[float, List[float], Tensor], batch: Optional[Tensor] = None, start: Optional[Union[float, List[float], Tensor]] = None, end: Optional[Union[float, List[float], Tensor]] = None, ) -> Tensor: r"""Voxel grid pooling from the, *e.g.*, `Dynamic Edge-Conditioned Filters in Convolutional Networks on Graphs <https://arxiv.org/abs/1704.02901>`_ paper, which overlays a regular grid of user-defined size over a point cloud and clusters all points within the same voxel. Args: pos (Tensor): Node position matrix :math:`\mathbf{X} \in \mathbb{R}^{(N_1 + \ldots + N_B) \times D}`. size (float or [float] or Tensor): Size of a voxel (in each dimension). batch (LongTensor, optional): Batch vector :math:`\mathbf{b} \in {\{ 0, \ldots,B-1\}}^N`, which assigns each node to a specific example. (default: :obj:`None`) start (float or [float] or Tensor, optional): Start coordinates of the grid (in each dimension). If set to :obj:`None`, will be set to the minimum coordinates found in :attr:`pos`. (default: :obj:`None`) end (float or [float] or Tensor, optional): End coordinates of the grid (in each dimension). If set to :obj:`None`, will be set to the maximum coordinates found in :attr:`pos`. (default: :obj:`None`) :rtype: :class:`LongTensor` """ if grid_cluster is None: raise ImportError('`voxel_grid` requires `torch-cluster`.') pos = pos.unsqueeze(-1) if pos.dim() == 1 else pos num_nodes, dim = pos.size() size = size.tolist() if torch.is_tensor(size) else size start = start.tolist() if torch.is_tensor(start) else start end = end.tolist() if torch.is_tensor(end) else end size, start, end = repeat(size, dim), repeat(start, dim), repeat(end, dim) if batch is None: batch = torch.zeros(pos.shape[0], dtype=torch.long) pos = torch.cat([pos, batch.unsqueeze(-1).type_as(pos)], dim=-1) size = size + [1] start = None if start is None else start + [0] end = None if end is None else end + [batch.max().item()] size = torch.tensor(size, dtype=pos.dtype, device=pos.device) if start is not None: start = torch.tensor(start, dtype=pos.dtype, device=pos.device) if end is not None: end = torch.tensor(end, dtype=pos.dtype, device=pos.device) return grid_cluster(pos, size, start, end)
def __init__( self, in_channels: Union[int, Tuple[int, int]], out_channels: int, dim: int, kernel_size: Union[int, List[int]], is_open_spline: bool = True, degree: int = 1, aggr: str = 'mean', root_weight: bool = True, bias: bool = True, **kwargs, ): super().__init__(aggr=aggr, **kwargs) if spline_basis is None: raise ImportError("'SplineConv' requires 'torch-spline-conv'") self.in_channels = in_channels self.out_channels = out_channels self.dim = dim self.degree = degree self.root_weight = root_weight kernel_size = torch.tensor(repeat(kernel_size, dim), dtype=torch.long) self.register_buffer('kernel_size', kernel_size) is_open_spline = repeat(is_open_spline, dim) is_open_spline = torch.tensor(is_open_spline, dtype=torch.uint8) self.register_buffer('is_open_spline', is_open_spline) if isinstance(in_channels, int): in_channels = (in_channels, in_channels) self.K = kernel_size.prod().item() if in_channels[0] > 0: self.weight = Parameter( torch.Tensor(self.K, in_channels[0], out_channels)) else: self.weight = torch.nn.parameter.UninitializedParameter() self._hook = self.register_forward_pre_hook( self.initialize_parameters) if root_weight: self.lin = Linear(in_channels[1], out_channels, bias=False, weight_initializer='uniform') if bias: self.bias = Parameter(torch.Tensor(out_channels)) else: self.register_parameter('bias', None) self.reset_parameters()
def __init__(self, in_channels, hidden_channels, out_channels, depth, pool_ratios=0.5, sum_res=True, act=F.relu): super(GraphUNet, self).__init__() assert depth >= 1 self.in_channels = in_channels self.hidden_channels = hidden_channels self.out_channels = out_channels self.depth = depth self.pool_ratios = repeat(pool_ratios, depth) self.act = act self.sum_res = sum_res channels = hidden_channels self.down_convs = torch.nn.ModuleList() self.pools = torch.nn.ModuleList() self.down_convs.append(conv(in_channels, channels, improved=True)) for i in range(depth): self.pools.append(TopKPooling(channels, self.pool_ratios[i])) self.down_convs.append(conv(channels, channels, improved=True)) in_channels = channels if sum_res else 2 * channels self.up_convs = torch.nn.ModuleList() for i in range(depth - 1): self.up_convs.append(conv(in_channels, channels, improved=True)) self.up_convs.append(conv(in_channels, out_channels, improved=True)) self.reset_parameters()
def __init__(self, data, size, num_hops, batch_size=1, shuffle=False, drop_last=False, bipartite=True, add_self_loops=False, flow='source_to_target'): self.data = data self.size = repeat(size, num_hops) self.num_hops = num_hops self.batch_size = batch_size self.shuffle = shuffle self.drop_last = drop_last self.bipartite = bipartite self.add_self_loops = add_self_loops self.flow = flow assert flow in ['source_to_target', 'target_to_source'] self.i, self.j = (0, 1) if flow == 'target_to_source' else (1, 0) self.edge_index_i, self.e_assoc = data.edge_index[self.i].sort() self.edge_index_j = data.edge_index[self.j, self.e_assoc] deg = degree(self.edge_index_i, data.num_nodes, dtype=torch.long) self.cumdeg = torch.cat([deg.new_zeros(1), deg.cumsum(0)]) self.tmp = torch.empty(data.num_nodes, dtype=torch.long)
def __init__(self, data, size, num_hops, batch_size=1, shuffle=False, drop_last=False, bipartite=True, add_self_loops=False, flow='source_to_target'): self.data = data self.size = repeat(size, num_hops) self.num_hops = num_hops self.batch_size = batch_size self.shuffle = shuffle self.drop_last = drop_last self.bipartite = bipartite self.add_self_loops = add_self_loops self.flow = flow self.edge_index = data.edge_index self.e_id = torch.arange(self.edge_index.size(1)) if bipartite and add_self_loops: tmp = segregate_self_loops(self.edge_index, self.e_id) self.edge_index, self.e_id, self.edge_index_loop = tmp[:3] self.e_id_loop = self.e_id.new_full((data.num_nodes, ), -1) self.e_id_loop[tmp[2][0]] = tmp[3] assert flow in ['source_to_target', 'target_to_source'] self.i, self.j = (0, 1) if flow == 'target_to_source' else (1, 0) edge_index_i, self.e_assoc = self.edge_index[self.i].sort() self.edge_index_j = self.edge_index[self.j, self.e_assoc] deg = degree(edge_index_i, data.num_nodes, dtype=torch.long) self.cumdeg = torch.cat([deg.new_zeros(1), deg.cumsum(0)]) self.tmp = torch.empty(data.num_nodes, dtype=torch.long)
def __init__( self, in_channels: int, hidden_channels: int, out_channels: int, depth: int = 1, pool_ratios: float = 0.5, act=F.relu, variational: bool = True, ): super(VariationalGraphEncoder, self).__init__() assert depth >= 1 self.depth = depth self.pool_ratios = repeat(pool_ratios, depth) self.act = act self.variational = variational self.down_convs = nn.ModuleList() self.pools = nn.ModuleList() self.down_convs.append(GCNConv(in_channels, hidden_channels, improved=True)) for i in range(depth): self.pools.append(TopKPooling(hidden_channels, self.pool_ratios[i])) self.down_convs.append( GCNConv(hidden_channels, hidden_channels, improved=True) ) self.conv_mu = GCNConv(hidden_channels, out_channels) if self.variational: self.conv_logstd = GCNConv(hidden_channels, out_channels) self.reset_parameters()
def __init__(self, in_channels, hidden_channels, out_channels, depth, params={}, pool_ratios=0.5, sum_res=True, act=F.relu): super(HillGraphUNet, self).__init__() assert depth >= 1 self.in_channels = in_channels self.hidden_channels = hidden_channels self.out_channels = out_channels self.depth = depth self.undirected_graphs = params['undirected_graphs'] self.n_attention_heads = params['n_attention_heads'] #self.attention_concat = params['attention_concat'] self.pool_ratios = repeat(pool_ratios, depth) self.act = act self.sum_res = sum_res self.use_batchnorm = params['use_batchnorm'] channels = hidden_channels self.down_convs = torch.nn.ModuleList() self.pools = torch.nn.ModuleList() self.down_bn = torch.nn.ModuleList() # batch norm after each GCN block # select between GCNConv and GATConv (graph attention) and BidirectionalGraphConv if params['bidirectional_graph_conv']: graph_conv = lambda in_c, out_c: BidirectionalGraphConv(in_c, out_c, n_attention_heads=self.n_attention_heads, concat=True) elif self.n_attention_heads == 0: graph_conv = lambda in_c, out_c: GCNBlock(in_c, channels, out_c) else: # we divide out_c by n_attention_heads because GATConv output is # by default the concatenation of all heads. # to get out_c channels at total, # we specify the attention heads with (out_c / n_attention_heads) output channels each # the user must verify that n_attention_heads divides out_c. n_heads Outputs of size outc//n_heads concatenated # one to each other >>> output of size out_c graph_conv = lambda in_c, out_c: GATConv(in_c, out_c // self.n_attention_heads, heads=self.n_attention_heads) self.down_convs.append(graph_conv(in_channels, channels)) self.down_bn.append(torch.nn.BatchNorm1d(channels)) for i in range(depth): self.pools.append(TopKPooling(channels, self.pool_ratios[i])) self.down_convs.append(graph_conv(channels, channels)) self.down_bn.append(torch.nn.BatchNorm1d(channels)) in_channels = channels if sum_res else 2 * channels self.up_convs = torch.nn.ModuleList() self.up_bn = torch.nn.ModuleList() for i in range(depth - 1): self.up_convs.append(graph_conv(in_channels, channels)) self.up_bn.append(torch.nn.BatchNorm1d(channels)) self.up_convs.append(graph_conv(in_channels, out_channels)) self.up_bn.append(torch.nn.BatchNorm1d(out_channels)) self.reset_parameters()
def voxel_grid(pos, batch, size, start=None, end=None): pos = pos.unsqueeze(-1) if pos.dim() == 1 else pos num_nodes, dim = pos.size() size, start, end = repeat(size, dim), repeat(start, dim), repeat(end, dim) pos = torch.cat([pos, batch.unsqueeze(-1).type_as(pos)], dim=-1) size = size + [1] start = None if start is None else start + [0] end = None if end is None else end + [batch.max().item()] size = torch.tensor(size, dtype=pos.dtype, device=pos.device) if start is not None: start = torch.tensor(start, dtype=pos.dtype, device=pos.device) if end is not None: end = torch.tensor(end, dtype=pos.dtype, device=pos.device) return grid_cluster(pos, size, start, end)
def __init__(self, in_channels: Union[int, Tuple[int, int]], out_channels: int, dim: int, kernel_size: Union[int, List[int]], is_open_spline: bool = True, degree: int = 1, aggr: str = 'mean', root_weight: bool = True, bias: bool = True, **kwargs): # yapf: disable super(SplineConv, self).__init__(aggr=aggr, **kwargs) if spline_basis is None: raise ImportError('`SplineConv` requires `torch-spline-conv`.') self.in_channels = in_channels self.out_channels = out_channels self.dim = dim self.degree = degree kernel_size = torch.tensor(repeat(kernel_size, dim), dtype=torch.long) self.register_buffer('kernel_size', kernel_size) is_open_spline = repeat(is_open_spline, dim) is_open_spline = torch.tensor(is_open_spline, dtype=torch.uint8) self.register_buffer('is_open_spline', is_open_spline) if isinstance(in_channels, int): in_channels = (in_channels, in_channels) K = kernel_size.prod().item() self.weight = Parameter(torch.Tensor(K, in_channels[0], out_channels)) if root_weight: self.root = Parameter(torch.Tensor(in_channels[1], out_channels)) else: self.register_parameter('root', None) if bias: self.bias = Parameter(torch.Tensor(out_channels)) else: self.register_parameter('bias', None) self.reset_parameters()
def __init__(self, in_channels, out_channels, dim, kernel_size, is_open_spline=True, degree=1, norm=True, root_weight=True, bias=True): super(SplineConv, self).__init__() self.in_channels = in_channels self.out_channels = out_channels self.degree = degree self.norm = norm kernel_size = torch.tensor(repeat(kernel_size, dim), dtype=torch.long) self.register_buffer('kernel_size', kernel_size) is_open_spline = repeat(is_open_spline, dim) is_open_spline = torch.tensor(is_open_spline, dtype=torch.uint8) self.register_buffer('is_open_spline', is_open_spline) K = kernel_size.prod().item() self.weight = Parameter(torch.Tensor(K, in_channels, out_channels)) if root_weight: self.root = Parameter(torch.Tensor(in_channels, out_channels)) else: self.register_parameter('root', None) if bias: self.bias = Parameter(torch.Tensor(out_channels)) else: self.register_parameter('bias', None) self.reset_parameters()
def __init__(self, in_channels, hidden_channels, out_channels, depth, pool_ratios=0.5, sum_res=True, act=F.relu, num_relations=4): super(GraphRUNet, self).__init__() assert depth >= 1 self.in_channels = in_channels self.hidden_channels = hidden_channels self.out_channels = out_channels self.depth = depth self.pool_ratios = repeat(pool_ratios, depth) self.act = act self.sum_res = sum_res channels = hidden_channels self.down_convs = torch.nn.ModuleList() self.pools = torch.nn.ModuleList() #self.down_convs.append(FastRGCNConv(in_channels, channels, num_relations=num_relations)) self.down_convs.append(GCNConv(in_channels, channels)) print("IN_CHANNELS", in_channels) #, num_relations=num_relations)) for i in range(depth): self.pools.append(TopKPooling(channels, self.pool_ratios[i])) self.down_convs.append(GCNConv(channels, channels)) in_channels = channels if sum_res else 2 * channels self.up_convs = torch.nn.ModuleList() for i in range(depth - 1): self.up_convs.append(GCNConv(channels, channels)) self.up_convs.append(GCNConv(channels, channels)) self.reset_parameters()
def test_repeat(): assert repeat(None, length=4) is None assert repeat(4, length=4) == [4, 4, 4, 4] assert repeat([2, 3, 4], length=4) == [2, 3, 4, 4] assert repeat([1, 2, 3, 4], length=4) == [1, 2, 3, 4] assert repeat([1, 2, 3, 4, 5], length=4) == [1, 2, 3, 4]