def __init__()

in source/sagemaker/sagemaker_graph_entity_resolution/dgl_entity_resolution/model.py [0:0]

• 19 lines of code
• 10 McCabe index (conditional complexity)

    def __init__(self, g, in_size, hidden_size, n_layers):
        super(HeteroRGCN, self).__init__()
        # Use trainable node embeddings as featureless inputs.
        embed_dict = {ntype: nn.Parameter(torch.Tensor(g.number_of_nodes(ntype), in_size['default']))
                      for ntype in g.ntypes if ntype != 'user' and ntype != 'website'}
        for key, embed in embed_dict.items():
            nn.init.xavier_uniform_(embed)
        self.embed = nn.ParameterDict(embed_dict)

        # Prepare R-GCN layer input output size for each relation
        in_sizes = []
        for srctype, etype, dsttype in g.canonical_etypes:
            if srctype in in_size:
                in_sizes.append(in_size[srctype])
            else:
                in_sizes.append(in_size['default'])

        hidden_sizes = [hidden_size] * len(g.etypes)

        # create layers
        layers = [HeteroRGCNLayer(in_sizes, hidden_sizes, g.etypes)]
        if n_layers > 1:
            # additional hidden layers
            for i in range(n_layers - 1):
                layers.append(HeteroRGCNLayer(hidden_sizes, hidden_sizes, g.etypes))
        self.layers = nn.Sequential(*layers)