def __init__(self, in_size, kernel_size, G_size=None, bias=False, normalize=False, activation='relu', dropout=0., b_norm=False, pooling='sum', pack_batch=False): # Accepted methods for pooling are avg, max and sum super().__init__() self.in_size = in_size self.out_size = kernel_size self.G_size = G_size self._pooling = get_pooling(pooling) self.pack_batch = pack_batch self.net = FCLayer(self.in_size, self.out_size, activation=activation, bias=bias, dropout=dropout, b_norm=b_norm) self.normalize = normalize self.reset_parameters()
def __init__(self, feat_dim, nedges, out_size=32, depth=2, pooling='max'): super(EdgeGraphLayer, self).__init__() self.feat_dim = feat_dim self.edge_dim = nedges self.out_size = out_size self.update_layers = nn.ModuleList() self.input_layer = nn.Linear(self.feat_dim, self.out_size) self.depth = depth for d in range(self.depth): self.update_layers.append( FCLayer(self.out_size + self.edge_dim, self.out_size)) self.pooling = get_pooling(pooling) self.readout = FCLayer(self.out_size, self.out_size)
def __init__(self, in_size, kernel_size, G_size=None, eps=None, net=None, pooling="sum", pack_batch=False, **kwargs): super(GINLayer, self).__init__() self.in_size = in_size self.out_size = kernel_size self._pooling = get_pooling(pooling) self.pack_batch = pack_batch self.G_size = G_size kwargs.pop("normalize", None) self.net = (net or FCLayer)(in_size, kernel_size, **kwargs) self.chosen_eps = eps if eps is None: self.eps = torch.nn.Parameter(torch.Tensor([0])) else: self.register_buffer('eps', torch.Tensor([eps])) self.reset_parameters()
def __init__(self, input_dim, out_dim, feat_dim=0, nedges=1, gather_dim=72, conv_dims=[64, 64], conv_dims_after=[128, 128], linear_dim=[32], dropout=0.2, gather="agg", pool_arch={}, pool_loss=False, **kwargs): super(Encoder, self).__init__() GraphLayer = kwargs.pop( "GraphConv", GLayer) # using GIN convolution when none is specified activation = kwargs.pop("activation", nn.LeakyReLU()) # Leaky relu self.input_dim = input_dim self.out_dim = out_dim self.feat_dim = feat_dim # when additional features are available for the node self.pool_loss = pool_loss self.nedges = nedges self.conv_layers1 = nn.ModuleList() self.conv_layers2 = nn.ModuleList() input_dim += feat_dim if nedges > 1: self.edge_layer = EdgeGraphLayer(input_dim, nedges, method='cat', b_norm=False, dropout=0.0, activation=activation) input_dim = self.edge_layer.output_dim + self.feat_dim for cdim in conv_dims: self.conv_layers1.append( GraphLayer(input_dim, kernel_size=cdim, b_norm=False, dropout=dropout, activation=activation)) input_dim = cdim self.pool_layer = get_graph_coarsener(**pool_arch) self.pool_arch = pool_arch.get("arch", "") if self.pool_layer is not None: self.pool_layer = self.pool_layer(input_dim) if self.pool_arch == "laplacian": input_dim = self.pool_layer.cluster_dim for cdim in conv_dims_after: self.conv_layers2.append( GraphLayer(input_dim, kernel_size=cdim, b_norm=False, dropout=dropout, activation=activation)) input_dim = cdim if gather == "agg": self.agg_layer = AggLayer(input_dim, gather_dim, dropout) elif gather in ["attn", "gated"]: self.agg_layer = get_pooling("attn", input_dim=input_dim, output_dim=gather_dim, dropout=dropout) else: gather_dim = input_dim self.agg_layer = get_pooling(gather) # multi dense layer input_dim = gather_dim self.linear_layers = nn.ModuleList() for ldim in linear_dim: self.linear_layers.append( FCLayer(input_dim, ldim, activation=activation, **kwargs)) input_dim = ldim self.output_layer = nn.Linear(input_dim, out_dim) self.pool_inspect = None self.pooling_loss = []