GPT2.cpp 3.96 KB
Newer Older
liucong's avatar
liucong committed
1
2
#include <GPT2.h>

3
4
#include <migraphx/onnx.hpp>
#include <migraphx/gpu/target.hpp>
liucong's avatar
liucong committed
5
6

#include <Filesystem.h>
7
8
9
10
11
12
13
14
#include <SimpleLog.h>
#include <algorithm>
#include <stdexcept>
#include <tokenization.h>

namespace migraphxSamples
{

liucong's avatar
liucong committed
15
GPT2::GPT2() {}
liucong's avatar
liucong committed
16

liucong's avatar
liucong committed
17
GPT2::~GPT2() {}
18

liucong's avatar
liucong committed
19
ErrorCode GPT2::Initialize()
20
{
liucong's avatar
liucong committed
21
    // 获取模型文件
liucong's avatar
liucong committed
22
    std::string modelPath = "../Resource/GPT2_shici.onnx";
23
24
25

    // 设置最大输入shape
    migraphx::onnx_options onnx_options;
liucong's avatar
liucong committed
26
    onnx_options.map_input_dims["input"] = {1, 1000};
27
28

    // 加载模型
liucong's avatar
liucong committed
29
    if(!Exists(modelPath))
liucong's avatar
liucong committed
30
    {
liucong's avatar
liucong committed
31
        LOG_ERROR(stdout, "%s not exist!\n", modelPath.c_str());
liucong's avatar
liucong committed
32
33
        return MODEL_NOT_EXIST;
    }
liucong's avatar
liucong committed
34
35
    net = migraphx::parse_onnx(modelPath, onnx_options);
    LOG_INFO(stdout, "succeed to load model: %s\n", GetFileName(modelPath).c_str());
36

liucong's avatar
liucong committed
37
    // 获取模型输入/输出节点信息
liucong's avatar
liucong committed
38
    std::unordered_map<std::string, migraphx::shape> inputs  = net.get_inputs();
liucong's avatar
liucong committed
39
    std::unordered_map<std::string, migraphx::shape> outputs = net.get_outputs();
liucong's avatar
liucong committed
40
41
    inputName                                                = inputs.begin()->first;
    inputShape                                               = inputs.begin()->second;
42
43
44
45
46
47

    // 设置模型为GPU模式
    migraphx::target gpuTarget = migraphx::gpu::target{};

    // 编译模型
    migraphx::compile_options options;
liucong's avatar
liucong committed
48
49
50
51
    options.device_id    = 0;    // 设置GPU设备,默认为0号设备
    options.offload_copy = true; // 设置offload_copy
    net.compile(gpuTarget, options);
    LOG_INFO(stdout, "succeed to compile model: %s\n", GetFileName(modelPath).c_str());
52
53
54
55

    return SUCCESS;
}

liucong's avatar
liucong committed
56
static bool CompareM(Predictions a, Predictions b) { return a.predictionvalue > b.predictionvalue; }
57

liucong's avatar
liucong committed
58
long unsigned int GPT2::Inference(const std::vector<long unsigned int>& input_id)
59
60
61
{

    long unsigned int input[1][input_id.size()];
liucong's avatar
liucong committed
62
    for(int j = 0; j < input_id.size(); ++j)
63
64
65
66
67
68
    {
        input[0][j] = input_id[j];
    }

    // 设置输入shape
    std::vector<std::vector<std::size_t>> inputShapes;
liucong's avatar
liucong committed
69
    inputShapes.push_back({1, input_id.size()});
70

liucong's avatar
liucong committed
71
    // 创建输入数据
liucong's avatar
liucong committed
72
    std::unordered_map<std::string, migraphx::argument> inputData;
liucong's avatar
liucong committed
73
74
    inputData[inputName] = migraphx::argument{migraphx::shape(inputShape.type(), inputShapes[0]),
                                              (long unsigned int*)input};
75
76
77
78
79

    // 推理
    std::vector<migraphx::argument> results = net.eval(inputData);

    // 获取输出节点的属性
liucong's avatar
liucong committed
80
81
82
83
    migraphx::argument result   = results[0];
    migraphx::shape outputShape = result.get_shape();     // 输出节点的shape
    int numberOfOutput          = outputShape.elements(); // 输出节点元素的个数
    float* data                 = (float*)result.data();  // 输出节点数据指针
84
85
86
87

    // 保存推理结果
    long unsigned int n = 0;
    std::vector<Predictions> resultsOfPredictions(22557);
liucong's avatar
liucong committed
88
    for(int i = (input_id.size() - 1) * 22557; i < input_id.size() * 22557; ++i)
89
    {
liucong's avatar
liucong committed
90
        resultsOfPredictions[n].index           = n;
91
92
93
94
95
96
97
98
99
100
101
102
103
104
        resultsOfPredictions[n].predictionvalue = data[i];
        ++n;
    }

    // 对于[UNK]的概率设为无穷小,模型的预测结果不可能是[UNK]
    resultsOfPredictions[100].predictionvalue = -10000;

    // 排序
    std::sort(resultsOfPredictions.begin(), resultsOfPredictions.end(), CompareM);

    return resultsOfPredictions[0].index;
}

ErrorCode GPT2::Preprocessing(cuBERT::FullTokenizer tokenizer,
liucong's avatar
liucong committed
105
106
                              char* question,
                              std::vector<long unsigned int>& input_id)
107
108
{
    // 分词操作
liucong's avatar
liucong committed
109
    int max_seq_length = 1000;
110
111
112
113
114
115
    std::vector<std::string> tokens_question;
    tokens_question.reserve(max_seq_length);
    tokenizer.tokenize(question, &tokens_question, max_seq_length);

    // 保存编码信息
    input_id.push_back(tokenizer.convert_token_to_id("[CLS]"));
liucong's avatar
liucong committed
116
    for(int i = 0; i < tokens_question.size(); ++i)
117
118
119
120
121
122
123
    {
        input_id.push_back(tokenizer.convert_token_to_id(tokens_question[i]));
    }

    return SUCCESS;
}

liucong's avatar
liucong committed
124
} // namespace migraphxSamples