def __init__()

in experiments/codes/model/gat/edge_gat.py [0:0]

• 106 lines of code
• 8 McCabe index (conditional complexity)

    def __init__(self, config, shared_embeddings=None):
        super(GatedGatEncoder, self).__init__(config)

        # flag to enable one-hot embedding if needed
        self.graph_mode = True
        self.one_hot = self.config.model.gat.emb_type == "one-hot"
        self.edgeConvs = []

        ## Add EdgeGATConv params
        for l in range(config.model.gat.num_layers):
            in_channels = config.model.relation_embedding_dim
            out_channels = config.model.relation_embedding_dim
            heads = config.model.gat.num_heads
            edge_dim = config.model.relation_embedding_dim

            weight = torch.Tensor(size=(in_channels, heads * out_channels)).to(
                config.general.device
            )
            weight.requires_grad = True
            self.add_weight(
                weight,
                "GatedEdgeGATConv.{}.weight".format(l),
                initializer=glorot,
                weight_norm=config.model.weight_norm,
            )

            att = torch.Tensor(size=(1, heads, 2 * out_channels + edge_dim)).to(
                config.general.device
            )
            att.requires_grad = True
            self.add_weight(
                att,
                "GatedEdgeGATConv.{}.att".format(l),
                initializer=glorot,
                weight_norm=config.model.weight_norm,
            )

            if l == 0:
                # only add the gru weights once
                gru_weight_ih = torch.Tensor(
                    size=(out_channels + edge_dim, 3 * out_channels)
                ).to(config.general.device)
                gru_weight_ih.requires_grad = True
                self.add_weight(
                    gru_weight_ih,
                    "GatedEdgeGATConv.{}.gru_w_ih".format("_all_"),
                    weight_norm=config.model.weight_norm,
                )

                gru_weight_hh = torch.Tensor(size=(out_channels, 3 * out_channels)).to(
                    config.general.device
                )
                gru_weight_hh.requires_grad = True
                self.add_weight(
                    gru_weight_hh,
                    "GatedEdgeGATConv.{}.gru_w_hh".format("_all_"),
                    weight_norm=config.model.weight_norm,
                )

                gru_bias_ih = torch.Tensor(size=(3 * out_channels,)).to(
                    config.general.device
                )
                gru_bias_ih.requires_grad = True
                self.add_weight(
                    gru_bias_ih,
                    "GatedEdgeGATConv.{}.gru_b_ih".format("_all_"),
                    initializer=(uniform, 1),
                    weight_norm=config.model.weight_norm,
                )

                gru_bias_hh = torch.Tensor(size=(3 * out_channels,)).to(
                    config.general.device
                )
                gru_bias_hh.requires_grad = True
                self.add_weight(
                    gru_bias_hh,
                    "GatedEdgeGATConv.{}.gru_b_hh".format("_all_"),
                    initializer=(uniform, 1),
                    weight_norm=config.model.weight_norm,
                )

            if config.model.gat.bias and config.model.gat.concat:
                bias = torch.Tensor(size=(heads * out_channels,)).to(
                    config.general.device
                )
                bias.requires_grad = True
                self.add_weight(
                    bias,
                    "GatedEdgeGATConv.{}.bias".format(l),
                    initializer=(uniform, 1),
                    weight_norm=config.model.weight_norm,
                )
            elif config.model.gat.bias and not config.model.gat.concat:
                bias = torch.Tensor(size=(out_channels,)).to(config.general.device)
                bias.requires_grad = True
                self.add_weight(
                    bias,
                    "GatedEdgeGATConv.{}.bias".format(l),
                    initializer=(uniform, 1),
                    weight_norm=config.model.weight_norm,
                )

            self.edgeConvs.append(
                GatedEdgeGatConv(
                    in_channels,
                    out_channels,
                    edge_dim,
                    heads=heads,
                    concat=config.model.gat.concat,
                    negative_slope=config.model.gat.negative_slope,
                    dropout=config.model.gat.dropout,
                    bias=config.model.gat.bias,
                )
            )

        ## Add classify params
        in_class_dim = (
            config.model.relation_embedding_dim * 2
            + config.model.relation_embedding_dim
        )
        self.add_classify_weights(in_class_dim)