Casrel模型pytorch版本

以下是CasRel模型的PyTorch版本示例代码：

import torch
import torch.nn as nn
import torch.nn.functional as F

class CasRel(nn.Module):
    def __init__(self, num_entities, num_relations, hidden_size):
        super(CasRel, self).__init__()
        self.num_entities = num_entities
        self.num_relations = num_relations
        self.hidden_size = hidden_size
        
        self.entity_embedding = nn.Embedding(num_entities, hidden_size)
        self.relation_embedding = nn.Embedding(num_relations, hidden_size)
        self.query_linear = nn.Linear(hidden_size, hidden_size)
        self.key_linear = nn.Linear(hidden_size, hidden_size)
        self.value_linear = nn.Linear(hidden_size, hidden_size)
        self.fc = nn.Linear(hidden_size, num_relations)
    
    def forward(self, entity_ids, relation_ids, entity_positions):
        batch_size = entity_ids.size(0)
        entity_embs = self.entity_embedding(entity_ids)
        relation_embs = self.relation_embedding(relation_ids)
        query = self.query_linear(entity_embs)
        key = self.key_linear(entity_embs)
        value = self.value_linear(entity_embs)
        attn_weights = torch.bmm(query, key.transpose(1,2))
        attn_weights = F.softmax(attn_weights, dim=2)
        attn_output = torch.bmm(attn_weights, value)
        entity_positions = entity_positions.unsqueeze(2)
        entity_positions = entity_positions.expand(batch_size, 1, self.hidden_size)
        entity_positions = entity_positions.float()
        entity_positions = entity_positions.cuda()
        entity_pos_embs = entity_positions * entity_embs
        output = torch.cat([attn_output, entity_pos_embs, relation_embs], dim=1)
        output = self.fc(output)
        return output

该模型包括实体嵌入层、关系嵌入层、查询线性层、键线性层、值线性层、多头注意力层和全连接层。在前向传递过程中，输入实体ID、关系ID和实体位置，经过实体嵌入、关系嵌入、多头注意力计算、实体位置嵌入和全连接层，最终输出预测的关系ID。