1. #include "Setting.h"
    2. #include "Random.h"
    3. #include "Reader.h"
    4. #include "Corrupt.h"
    5. #include "Test.h"
    6. #include <cstdlib>
    7. #include <pthread.h>
    9. extern "C" __declspec(dllexport)
    10. void setInPath(char *path);
    12. extern "C" __declspec(dllexport)
    13. void setTrainPath(char *path);
    15. extern "C" __declspec(dllexport)
    16. void setValidPath(char *path);
    18. extern "C" __declspec(dllexport)
    19. void setTestPath(char *path);
    21. extern "C" __declspec(dllexport)
    22. void setEntPath(char *path);
    24. extern "C" __declspec(dllexport)
    25. void setRelPath(char *path);
    27. extern "C" __declspec(dllexport)
    28. void setOutPath(char *path);
    30. extern "C" __declspec(dllexport)
    31. void setWorkThreads(INT threads);
    33. extern "C" __declspec(dllexport)
    34. void setBern(INT con);
    36. extern "C" __declspec(dllexport)
    37. INT getWorkThreads();
    39. extern "C" __declspec(dllexport)
    40. INT getEntityTotal();
    42. extern "C" __declspec(dllexport)
    43. INT getRelationTotal();
    45. extern "C" __declspec(dllexport)
    46. INT getTripleTotal();
    48. extern "C" __declspec(dllexport)
    49. INT getTrainTotal();
    51. extern "C" __declspec(dllexport)
    52. INT getTestTotal();
    54. extern "C" __declspec(dllexport)
    55. INT getValidTotal();
    57. extern "C" __declspec(dllexport)
    58. void randReset();
    60. extern "C" __declspec(dllexport)
    61. void importTrainFiles();
    63. struct Parameter {
    64.     INT id;
    65.     INT *batch_h;
    66.     INT *batch_t;
    67.     INT *batch_r;
    68.     REAL *batch_y;
    69.     INT batchSize;
    70.     INT negRate;
    71.     INT negRelRate;
    72.     bool p;
    73.     bool val_loss;
    74.     INT mode;
    75.     bool filter_flag;
    76. };
    78. void* getBatch(void* con) {
    79.     // 将传入的参数直接变成结构体形式进行接收
    80.     Parameter *para = (Parameter *)(con);
    81.     // 将para中对应的值放到局部变量中
    82.     INT id = para -> id;
    83.     INT *batch_h = para -> batch_h;
    84.     INT *batch_t = para -> batch_t;
    85.     INT *batch_r = para -> batch_r;
    86.     REAL *batch_y = para -> batch_y;
    87.     INT batchSize = para -> batchSize;
    88.     INT negRate = para -> negRate;
    89.     INT negRelRate = para -> negRelRate;
    90.     bool p = para -> p;
    91.     // 判断模式是否为训练
    92.     bool val_loss = para -> val_loss;
    93.     INT mode = para -> mode;
    94.     bool filter_flag = para -> filter_flag;
    96.     // 既然要并行计算嘛，那么不同线程处理的数据范围肯定是不一样的，这里就是划分了一下每个线程的操作范围
    97.     INT lef, rig;
    98.     if (batchSize % workThreads == 0) {
    99.         lef = id * (batchSize / workThreads);
    100.         rig = (id + 1) * (batchSize / workThreads);
    101.     } else {
    102.         lef = id * (batchSize / workThreads + 1);    
    103.         rig = (id + 1) * (batchSize / workThreads + 1);
    104.         if (rig > batchSize) rig = batchSize;
    105.     }
    107.     // 一个阈值，决定负采样头部与负采样尾部的比率，500表示各自一半
    108.     REAL prob = 500;
    109.     // 开始采样负例三元组，用于训练
    110.     if (val_loss == false) {
    111.         for (INT batch = lef; batch < rig; batch++) {
    112.             // rand的实现是赋予每个线程一个种子，这样每个线程可以生成自己的随机数（可能是因为所有线程都拿当前时间作为seed的话可能会造成重复？）
    113.             INT i = rand_max(id, trainTotal);
    114.             // 得到了处于0~三元组总数之间的一个随机id
    115.             batch_h[batch] = trainList[i].h;
    116.             batch_t[batch] = trainList[i].t;
    117.             batch_r[batch] = trainList[i].r;
    118.             batch_y[batch] = 1;
    119.             INT last = batchSize;
    120.             // 负采样实体negRate次
    121.             for (INT times = 0; times < negRate; times ++) {
    122.                 // 模式为普通采样，即负采样头部与负采样尾部各占一半
    123.                 if (mode == 0){
    124.                     if (bernFlag)
    125.                         prob = 1000 * right_mean[trainList[i].r] / (right_mean[trainList[i].r] + left_mean[trainList[i].r]);
    126.                     if (randd(id) % 1000 < prob) {
    127.                         batch_h[batch + last] = trainList[i].h;
    128.                         batch_t[batch + last] = corrupt_head(id, trainList[i].h, trainList[i].r);
    129.                         batch_r[batch + last] = trainList[i].r;
    130.                     } else {
    131.                         batch_h[batch + last] = corrupt_tail(id, trainList[i].t, trainList[i].r);
    132.                         batch_t[batch + last] = trainList[i].t;
    133.                         batch_r[batch + last] = trainList[i].r;
    134.                     }
    135.                     batch_y[batch + last] = -1;
    136.                     last += batchSize;
    137.                 }
    138.                 // 模式为头部采样，即仅负采样头部
    139.                 else {
    140.                     if(mode == -1){
    141.                         batch_h[batch + last] = corrupt_tail(id, trainList[i].t, trainList[i].r);
    142.                         batch_t[batch + last] = trainList[i].t;
    143.                         batch_r[batch + last] = trainList[i].r;
    144.                     }
    145.                     // 模式为尾部采样，即仅负采样尾部
    146.                     else {
    147.                         batch_h[batch + last] = trainList[i].h;
    148.                         batch_t[batch + last] = corrupt_head(id, trainList[i].h, trainList[i].r);
    149.                         batch_r[batch + last] = trainList[i].r;
    150.                     }
    151.                     batch_y[batch + last] = -1;
    152.                     last += batchSize;
    153.                 }
    154.             }
    155.             // 负采样关系negRelRate次
    156.             for (INT times = 0; times < negRelRate; times++) {
    157.                 batch_h[batch + last] = trainList[i].h;
    158.                 batch_t[batch + last] = trainList[i].t;
    159.                 batch_r[batch + last] = corrupt_rel(id, trainList[i].h, trainList[i].t, trainList[i].r, p);
    160.                 batch_y[batch + last] = -1;
    161.                 last += batchSize;
    162.             }
    163.         }
    164.     }
    165.     else
    166.     {
    167.         for (INT batch = lef; batch < rig; batch++)
    168.         {
    169.             batch_h[batch] = validList[batch].h;
    170.             batch_t[batch] = validList[batch].t;
    171.             batch_r[batch] = validList[batch].r;
    172.             batch_y[batch] = 1;
    173.         }
    174.     }
    175.     pthread_exit(NULL);
    176.     return ((void*)0);
    177. }
    179. extern "C" __declspec(dllexport)
    180. void sampling(
    181.         INT *batch_h, 
    182.         INT *batch_t, 
    183.         INT *batch_r, 
    184.         REAL *batch_y, 
    185.         INT batchSize, 
    186.         INT negRate = 1, 
    187.         INT negRelRate = 0, 
    188.         // mode=0代表普通采样，mode=1代表尾部采样，mode=-1代表头部采样
    189.         INT mode = 0,
    190.         bool filter_flag = true,
    191.         bool p = false, 
    192.         // val_loss应该代表的是是否为训练模式
    193.         bool val_loss = false
    194. ) {
    195.     // 根据线程数量，向内存分配指定的大小
    196.     pthread_t *pt = (pthread_t *)malloc(workThreads * sizeof(pthread_t));
    198.     // 根据线程数量，以及Parameter结构体的大小，向内存分配指定的大小
    199.     Parameter *para = (Parameter *)malloc(workThreads * sizeof(Parameter));
    201.     //初始化para结构体
    202.     for (INT threads = 0; threads < workThreads; threads++) {
    203.         para[threads].id = threads;
    204.         para[threads].batch_h = batch_h;
    205.         para[threads].batch_t = batch_t;
    206.         para[threads].batch_r = batch_r;
    207.         para[threads].batch_y = batch_y;
    208.         para[threads].batchSize = batchSize;
    209.         para[threads].negRate = negRate;
    210.         para[threads].negRelRate = negRelRate;
    211.         para[threads].p = p;
    212.         para[threads].val_loss = val_loss;
    213.         para[threads].mode = mode;
    214.         para[threads].filter_flag = filter_flag;
    215.         /*
    216.         创建线程
    217.         int pthread_create(
    218.              pthread_t *restrict tidp,   //新创建的线程ID指向的内存单元。
    219.              const pthread_attr_t *restrict attr,  //线程属性，默认为NULL
    220.              void *(*start_rtn)(void *), //新创建的线程从start_rtn函数的地址开始运行
    221.              void *restrict arg //默认为NULL。若上述函数需要参数，将参数放入结构中并将地址作为arg传入。
    222.               );
    223.         */
    224.         pthread_create(&pt[threads], NULL, getBatch, (void*)(para+threads));
    225.     }
    226.     // 收工
    227.     for (INT threads = 0; threads < workThreads; threads++)
    228.         pthread_join(pt[threads], NULL);
    230.     free(pt);
    231.     free(para);
    232. }
    234. int main() {
    235.     importTrainFiles();
    236.     return 0;
    237. }