1. # coding:utf-8
    2. import os
    3. import ctypes
    4. import numpy as np
    6. class TrainDataSampler(object):
    8.     def __init__(self, nbatches, datasampler):
    9.         self.nbatches = nbatches
    10.         self.datasampler = datasampler
    11.         self.batch = 0
    13.     def __iter__(self):
    14.         return self
    16.     def __next__(self):
    17.         self.batch += 1 
    18.         if self.batch > self.nbatches:
    19.             raise StopIteration()
    20.         return self.datasampler()
    22.     def __len__(self):
    23.         return self.nbatches
    25. class TrainDataLoader(object):
    27.     def __init__(self, 
    28.         in_path = "./",
    29.         tri_file = None,  # 三元组文件
    30.         ent_file = None,  # 实体文件记录实体与对应id之间的映射关系
    31.         rel_file = None,  # 关系文件记录关系与对应id之间的映射关系
    32.         batch_size = None,
    33.         nbatches = None,  # 根据三元组的总量计算有多少个batch
    34.         threads = 8,  # 读取数据的线程数量
    35.         sampling_mode = "normal",  # 采样方法，分为普通采样及交叉采样
    36.         bern_flag = False,  # 负采样时进行伯努利采样还是均匀采样
    37.         filter_flag = True,  # 进行负采样时是否要过滤掉存在于训练集中的三元组
    38.         neg_ent = 1,  # 负采样时一个正样本生成负样本的数量，默认情况一个正样本只生成一个负样本
    39.         neg_rel = 0):  # 负采样时一个关系生成负样本的数量
    41.         base_file = os.path.abspath(os.path.join(os.path.dirname(__file__), "../release/Base.dll"))
    42.         self.lib = ctypes.cdll.LoadLibrary(base_file)
    43.         """argtypes"""
    44.         self.lib.sampling.argtypes = [  # C与Python数据类型的转换，回调
    45.             ctypes.c_void_p,
    46.             ctypes.c_void_p,
    47.             ctypes.c_void_p,
    48.             ctypes.c_void_p,
    49.             ctypes.c_int64,
    50.             ctypes.c_int64,
    51.             ctypes.c_int64,
    52.             ctypes.c_int64,
    53.             ctypes.c_int64,
    54.             ctypes.c_int64,
    55.             ctypes.c_int64
    56.         ]
    57.         self.in_path = in_path
    58.         self.tri_file = tri_file
    59.         self.ent_file = ent_file
    60.         self.rel_file = rel_file
    61.         if in_path != None:
    62.             self.tri_file = in_path + "train2id.txt"
    63.             self.ent_file = in_path + "entity2id.txt"
    64.             self.rel_file = in_path + "relation2id.txt"
    65.         """set essential parameters"""
    66.         self.work_threads = threads
    67.         self.nbatches = nbatches
    68.         self.batch_size = batch_size
    69.         self.bern = bern_flag
    70.         self.filter = filter_flag
    71.         self.negative_ent = neg_ent  # 生成负例实体数量
    72.         self.negative_rel = neg_rel  # 生成负例关系数量
    73.         self.sampling_mode = sampling_mode
    74.         self.cross_sampling_flag = 0
    75.         self.read()
    77.     def read(self):
    78.         # 将文件路径之类的信息传入动态链接库
    79.         if self.in_path != None:
    80.             self.lib.setInPath(ctypes.create_string_buffer(self.in_path.encode(), len(self.in_path) * 2))
    81.         else:
    82.             self.lib.setTrainPath(ctypes.create_string_buffer(self.tri_file.encode(), len(self.tri_file) * 2))
    83.             self.lib.setEntPath(ctypes.create_string_buffer(self.ent_file.encode(), len(self.ent_file) * 2))
    84.             self.lib.setRelPath(ctypes.create_string_buffer(self.rel_file.encode(), len(self.rel_file) * 2))
    86.         # lib.method均表示调用动态链接库中的函数
    88.         # 设置是否进行伯努利采样
    89.         self.lib.setBern(self.bern)
    90.         # 设置工作线程
    91.         self.lib.setWorkThreads(self.work_threads)
    92.         # 重置线程的随机种子
    93.         self.lib.randReset()
    94.         # 底层c++开始将数据读入内存
    95.         self.lib.importTrainFiles()
    96.         # 获取实体、关系以及训练三元组的总数
    97.         self.relTotal = self.lib.getRelationTotal()
    98.         self.entTotal = self.lib.getEntityTotal()
    99.         self.tripleTotal = self.lib.getTrainTotal()
    101.         if self.batch_size == None:
    102.             self.batch_size = self.tripleTotal // self.nbatches
    103.         if self.nbatches == None:
    104.             self.nbatches = self.tripleTotal // self.batch_size
    106.         # batch_seq_size是每次训练时一个batch的数据，其中包括原始及负采样之后得到的三元组，所以在长度上要乘以负采样实体和关系的次数
    107.         self.batch_seq_size = self.batch_size * (1 + self.negative_ent + self.negative_rel)
    109.         # 初始化一个batch的数据，包含头实体、尾实体、关系、标签，其中batch_y为标签，1表示原始三元组，-1表示负采样替换后的三元组
    110.         self.batch_h = np.zeros(self.batch_seq_size, dtype=np.int64)
    111.         self.batch_t = np.zeros(self.batch_seq_size, dtype=np.int64)
    112.         self.batch_r = np.zeros(self.batch_seq_size, dtype=np.int64)
    113.         self.batch_y = np.zeros(self.batch_seq_size, dtype=np.float32)
    115.         # 获得tensor的地址，用于将地址传给c++，让底层代码完成数据的写入工作。这样就完成了数据的传递（居然还能这样）
    116.         self.batch_h_addr = self.batch_h.__array_interface__["data"][0]
    117.         self.batch_t_addr = self.batch_t.__array_interface__["data"][0]
    118.         self.batch_r_addr = self.batch_r.__array_interface__["data"][0]
    119.         self.batch_y_addr = self.batch_y.__array_interface__["data"][0]
    121.     # 普通采样
    122.     def sampling(self):
    123.         # 这里通过刚才获取的tensor地址就可以传给c++进行采样操作了
    124.         self.lib.sampling(
    125.             self.batch_h_addr,
    126.             self.batch_t_addr,
    127.             self.batch_r_addr,
    128.             self.batch_y_addr,
    129.             self.batch_size,
    130.             self.negative_ent,
    131.             self.negative_rel,
    132.             0,
    133.             self.filter,
    134.             0,
    135.             0
    136.         )
    137.         return {
    138.             "batch_h": self.batch_h, 
    139.             "batch_t": self.batch_t, 
    140.             "batch_r": self.batch_r, 
    141.             "batch_y": self.batch_y,
    142.             "mode": "normal"
    143.         }
    145.     def sampling_head(self):
    146.         self.lib.sampling(
    147.             self.batch_h_addr,
    148.             self.batch_t_addr,
    149.             self.batch_r_addr,
    150.             self.batch_y_addr,
    151.             self.batch_size,
    152.             self.negative_ent,
    153.             self.negative_rel,
    154.             -1,
    155.             self.filter,
    156.             0,
    157.             0
    158.         )
    159.         return {
    160.             "batch_h": self.batch_h,
    161.             "batch_t": self.batch_t[:self.batch_size],
    162.             "batch_r": self.batch_r[:self.batch_size],
    163.             "batch_y": self.batch_y,
    164.             "mode": "head_batch"
    165.         }
    167.     def sampling_tail(self):
    168.         self.lib.sampling(
    169.             self.batch_h_addr,
    170.             self.batch_t_addr,
    171.             self.batch_r_addr,
    172.             self.batch_y_addr,
    173.             self.batch_size,
    174.             self.negative_ent,
    175.             self.negative_rel,
    176.             1,
    177.             self.filter,
    178.             0,
    179.             0
    180.         )
    181.         return {
    182.             "batch_h": self.batch_h[:self.batch_size],
    183.             "batch_t": self.batch_t,
    184.             "batch_r": self.batch_r[:self.batch_size],
    185.             "batch_y": self.batch_y,
    186.             "mode": "tail_batch"
    187.         }
    189.     # 交叉采样，每遍历到一个batch交替进行头部采样以及尾部采样
    190.     def cross_sampling(self):
    191.         self.cross_sampling_flag = 1 - self.cross_sampling_flag 
    192.         if self.cross_sampling_flag == 0:
    193.             return self.sampling_head()
    194.         else:
    195.             return self.sampling_tail()
    197.     """interfaces to set essential parameters"""
    199.     def set_work_threads(self, work_threads):
    200.         self.work_threads = work_threads
    202.     def set_in_path(self, in_path):
    203.         self.in_path = in_path
    205.     def set_nbatches(self, nbatches):
    206.         self.nbatches = nbatches
    208.     def set_batch_size(self, batch_size):
    209.         self.batch_size = batch_size
    210.         self.nbatches = self.tripleTotal // self.batch_size
    212.     def set_ent_neg_rate(self, rate):
    213.         self.negative_ent = rate
    215.     def set_rel_neg_rate(self, rate):
    216.         self.negative_rel = rate
    218.     def set_bern_flag(self, bern):
    219.         self.bern = bern
    221.     def set_filter_flag(self, filter):
    222.         self.filter = filter
    224.     """interfaces to get essential parameters"""
    226.     def get_batch_size(self):
    227.         return self.batch_size
    229.     def get_ent_tot(self):
    230.         return self.entTotal
    232.     def get_rel_tot(self):
    233.         return self.relTotal
    235.     def get_triple_tot(self):
    236.         return self.tripleTotal
    238.     def __iter__(self):
    239.         if self.sampling_mode == "normal":
    240.             return TrainDataSampler(self.nbatches, self.sampling)
    241.         else:
    242.             return TrainDataSampler(self.nbatches, self.cross_sampling)
    244.     def __len__(self):
    245.         return self.nbatches