OpenVINO によるナレッジグラフの最適化#

# 事前学習済みモデルのチェックポイントへのパス 
modelpath = Path("models/conve.pt") 

# エンティティーと関係の埋め込み次元 
EMB_DIM = 300 

# 予測から考慮すべき上位 K 値 
TOP_K = 2 

# OpenVINO 変換に必要 
output_dir = Path("models") 
base_model_name = "conve" 

output_dir.mkdir(exist_ok=True) 

# PyTorch および OpenVINO IR モデルが保存されるパス 
ir_path = Path(output_dir / base_model_name).with_suffix(".xml")

data_folder = "data" 

# エンティティーとエンティティー ID を含むファイルをダウンロード 
entdatapath = download_file( 

"https://storage.openvinotoolkit.org/repositories/openvino_notebooks/data/data/text/countries_S1/kg_training_entids.txt", 
    directory=data_folder, 
) 

# 関係と関係 ID を含むファイルをダウンロード 
reldatapath = download_file( 

"https://storage.openvinotoolkit.org/repositories/openvino_notebooks/data/data/text/countries_S1/kg_training_relids.txt", 
    directory=data_folder, 
) 

# テスト・データ・ファイルをダウンロード 
testdatapath = download_file( 

"https://storage.openvinotoolkit.org/repositories/openvino_notebooks/data/data/json/countries_S1/e1rel_to_e2_ranking_test.json", 
    directory=data_folder, 
)

model_url = "https://storage.openvinotoolkit.org/repositories/openvino_notebooks/models/knowledge-graph-embeddings/conve.pt" 

download_file(model_url, filename=modelpath.name, directory=modelpath.parent)

# モデル実装リファレンス: https://github.com/TimDettmers/ConvE 
class ConvE(torch.nn.Module): 
    def __init__(self, num_entities, num_relations, emb_dim): 
        super(ConvE, self).__init__() 
        # y 次元に num_uniq_ent、x 次元に emb_dim を持つエンティティーと関係テーブルを埋め込む 
        self.emb_e = torch.nn.Embedding(num_entities, emb_dim, padding_idx=0) 
        self.ent_weights_matrix = torch.ones([num_entities, emb_dim], dtype=torch.float64) 
        self.emb_rel = torch.nn.Embedding(num_relations, emb_dim, padding_idx=0) 
        self.ne = num_entities 
        self.nr = num_relations 
        self.inp_drop = torch.nn.Dropout(0.2) 
        self.hidden_drop = torch.nn.Dropout(0.3) 
        self.feature_map_drop = torch.nn.Dropout2d(0.2) 
        self.loss = torch.nn.BCELoss() 
        self.conv1 = torch.nn.Conv2d(1, 32, (3, 3), 1, 0, bias=True) 
        self.bn0 = torch.nn.BatchNorm2d(1) 
        self.bn1 = torch.nn.BatchNorm2d(32) 
        self.ln0 = torch.nn.LayerNorm(emb_dim) 
        self.register_parameter("b", Parameter(torch.zeros(num_entities))) 
        self.fc = torch.nn.Linear(16128, emb_dim) 

    def init(self): 
        """Initializes the model""" 
        # Xavier の初期化 
        xavier_normal_(self.emb_e.weight.data) 
        xavier_normal_(self.emb_rel.weight.data) 

def forward(self, e1, rel): 
        """Forward pass on the model.
        :param e1: source entity 
        :param rel: relation between the source and target entities 
        Returns the model predictions for the target entities 
        """ 
        e1_embedded = self.emb_e(e1).view(-1, 1, 10, 30) 
        rel_embedded = self.emb_rel(rel).view(-1, 1, 10, 30) 
        stacked_inputs = torch.cat([e1_embedded, rel_embedded], 2) 
        stacked_inputs = self.bn0(stacked_inputs) 
        x = self.inp_drop(stacked_inputs) 
        x = self.conv1(x) 
        x = self.bn1(x) 
        x = F.relu(x) 
        x = self.feature_map_drop(x) 
        x = x.view(1, -1) 
        x = self.fc(x) 
        x = self.hidden_drop(x) 
        x = self.ln0(x) 
        x = F.relu(x) 
        x = torch.mm(x, self.emb_e.weight.transpose(1, 0)) 
        x = self.hidden_drop(x) 
        x += self.b.expand_as(x) 
        pred = torch.nn.functional.softmax(x, dim=1) 
        return pred

class DataLoader: 
    def __init__(self): 
        super(DataLoader, self).__init__() 

        self.ent_path = entdatapath 
        self.rel_path = reldatapath 
        self.test_file = testdatapath 
        self.entity_ids, self.ids2entities = self.load_data(data_path=self.ent_path) 
        self.rel_ids, self.ids2rel = self.load_data(data_path=self.rel_path) 
        self.test_triples_list = self.convert_triples(data_path=self.test_file) 

    def load_data(self, data_path): 
        """Creates a dictionary of data items with corresponding ids""" 
        item_dict, ids_dict = {}, {} 
        fp = open(data_path, "r") 
        lines = fp.readlines() 
        for line in lines: 
            name, id = line.strip().split("\t") 
            item_dict[name] = int(id) 
            ids_dict[int(id)] = name 
        fp.close() 
        return item_dict, ids_dict 

    def convert_triples(self, data_path): 
        """Creates a triple of source entity, relation and target entities""" 
        triples_list = [] 
        dp = open(data_path, "r") 
        lines = dp.readlines() 
        for line in lines: 
            item_dict = json.loads(line.strip()) 
            h = item_dict["e1"] 
            r = item_dict["rel"] 
            t = item_dict["e2_multi1"].split("\t") 
            hrt_list = [] 
            hrt_list.append(self.entity_ids[h]) 
            hrt_list.append(self.rel_ids[r]) 
            t_ents = [] 
            for t_idx in t: 
                t_ents.append(self.entity_ids[t_idx]) 
            hrt_list.append(t_ents) 
            triples_list.append(hrt_list) 
        dp.close() 
        return triples_list

data = DataLoader() 
num_entities = len(data.entity_ids) 
num_relations = len(data.rel_ids) 

model = ConvE(num_entities=num_entities, num_relations=num_relations, emb_dim=EMB_DIM) 
model.load_state_dict(torch.load(modelpath)) 
model.eval() 

pt_inf_times = [] 

triples_list = data.test_triples_list 
num_test_samples = len(triples_list) 
pt_acc = 0.0 
for i in range(num_test_samples): 
    test_sample = triples_list[i] 
    h, r, t = test_sample 
    start_time = time.time() 
    logits = model.forward(e1=torch.tensor(h), rel=torch.tensor(r)) 
    end_time = time.time() 
    pt_inf_times.append(end_time - start_time) 
    score, pred = torch.topk(logits, TOP_K, 1) 

    gt = np.array(sorted(t)) 
    pred = np.array(sorted(pred[0].cpu().detach())) 
    pt_acc += accuracy_score(gt, pred) 

avg_pt_time = np.mean(pt_inf_times) * 1000 
print(f"Average time taken for inference: {avg_pt_time} ms") 
print(f"Mean accuracy of the model on the test dataset: {pt_acc/num_test_samples}")

entitynames_dict = data.ids2entities 

ent = "san_marino" 
rel = "locatedin" 

h_idx = data.entity_ids[ent] 
r_idx = data.rel_ids[rel] 

logits = model.forward(torch.tensor(h_idx), torch.tensor(r_idx)) 
score, pred = torch.topk(logits, TOP_K, 1) 

for j, id in enumerate(pred[0].cpu().detach().numpy()): 
    pred_entity = entitynames_dict[id] 
    print(f"Source Entity: {ent}, Relation: {rel}, Target entity prediction: {pred_entity}")

print("Converting the trained conve model to IR format") 

ov_model = ov.convert_model(model, example_input=(torch.tensor(1), torch.tensor(1))) 
ov.save_model(ov_model, ir_path)

core = ov.Core() 
ov_model = core.read_model(model=ir_path)

# ゼロ除算を防止 
delimiter = max(avg_ov_time, np.finfo(float).eps) 

print(f"Speedup with OpenVINO optimizations: {round(float(avg_pt_time)/float(delimiter),2)} X")

print("Benchmark OpenVINO model using the benchmark app") 
! benchmark_app -m $ir_path -d $device.value -api async -t 10 -shape "input.1[1],input.2[1]"