Exemplo n.º 1
0
    def __init__(self,
                 device,
                 num_nodes,
                 h_dim,
                 out_dim,
                 num_rels,
                 num_bases=None,
                 num_hidden_layers=1,
                 dropout=0,
                 use_self_loop=False,
                 low_mem=False,
                 layer_norm=False):
        super(EntityClassify, self).__init__()
        self.device = th.device(device if device >= 0 else 'cpu')
        self.num_nodes = num_nodes
        self.h_dim = h_dim
        self.out_dim = out_dim
        self.num_rels = num_rels
        self.num_bases = None if num_bases < 0 else num_bases
        self.num_hidden_layers = num_hidden_layers
        self.dropout = dropout
        self.use_self_loop = use_self_loop
        self.low_mem = low_mem
        self.layer_norm = layer_norm

        self.layers = nn.ModuleList()
        # i2h
        self.layers.append(
            RelGraphConv(self.h_dim,
                         self.h_dim,
                         self.num_rels,
                         "basis",
                         self.num_bases,
                         activation=F.relu,
                         self_loop=self.use_self_loop,
                         low_mem=self.low_mem,
                         dropout=self.dropout))
        # h2h
        for idx in range(self.num_hidden_layers):
            self.layers.append(
                RelGraphConv(self.h_dim,
                             self.h_dim,
                             self.num_rels,
                             "basis",
                             self.num_bases,
                             activation=F.relu,
                             self_loop=self.use_self_loop,
                             low_mem=self.low_mem,
                             dropout=self.dropout))
        # h2o
        self.layers.append(
            RelGraphConv(self.h_dim,
                         self.out_dim,
                         self.num_rels,
                         "basis",
                         self.num_bases,
                         activation=None,
                         self_loop=self.use_self_loop,
                         low_mem=self.low_mem))
Exemplo n.º 2
0
 def __init__(self,
              num_nodes,
              hidden_dim,
              num_rels,
              num_layers=2,
              regularizer='basis',
              num_bases=None,
              dropout=0.0):
     """R-GCN连接预测模型
     
     :param num_nodes: int 顶点(实体)数
     :param hidden_dim: int 隐含特征维数
     :param num_rels: int 关系个数
     :param num_layers: int, optional R-GCN层数,默认为2
     :param regularizer: str, 'basis'/'bdd' 权重正则化方法,默认为'basis'
     :param num_bases: int, optional 基的个数,默认使用关系个数
     :param dropout: float, optional Dropout概率,默认为0
     """
     super().__init__()
     self.embed = nn.Embedding(num_nodes, hidden_dim)
     self.layers = nn.ModuleList()
     for i in range(num_layers):
         self.layers.append(
             RelGraphConv(hidden_dim,
                          hidden_dim,
                          num_rels,
                          regularizer,
                          num_bases,
                          activation=F.relu if i < num_layers - 1 else None,
                          self_loop=True,
                          low_mem=True,
                          dropout=dropout))
     self.score = DistMult(num_rels, hidden_dim)
Exemplo n.º 3
0
    def __init__(self, in_feat, hid_feat, out_feat, n_hidden_layers, num_rel):
        super().__init__()
        self.h_embedding = nn.Embedding(in_feat, hid_feat)
        self.mlp = MLPPredictor(hid_feat, out_feat)

        self.layers = nn.ModuleList([
            RelGraphConv(hid_feat,
                         hid_feat,
                         num_rel,
                         low_mem=True,
                         dropout=0.4) for _ in range(n_hidden_layers)
        ])
Exemplo n.º 4
0
    def __init__(self, in_dim, hid_dim, out_dim, n_layers, num_rel, dropout):
        super().__init__()
        self.in_dim = in_dim
        self.hid_dim = hid_dim
        self.out_dim = out_dim
        self.n_layers = n_layers
        self.num_rel = num_rel
        self.dropout = dropout

        self.h_embedding = nn.Embedding(in_dim, hid_dim)
        self.layers = nn.ModuleList([
            RelGraphConv(hid_dim,
                         hid_dim,
                         num_rel,
                         regularizer='basis',
                         num_bases=None,
                         low_mem=True,
                         dropout=self.dropout) for _ in range(n_layers)
        ])

        self.bilin_score = BiLinearPredictor(num_rel, hid_dim)