1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
|
#pragma once
#include "cnpy/cnpy.h"
#include "mblas/matrix.h"
namespace amunmt {
namespace CPU {
class NpzConverter {
private:
class NpyMatrixWrapper {
public:
NpyMatrixWrapper(const cnpy::NpyArray& npy)
: npy_(npy) {}
size_t size() const {
return size1() * size2();
}
float* data() const {
return (float*)npy_.data;
}
float operator()(size_t i, size_t j) const {
return ((float*)npy_.data)[i * size2() + j];
}
size_t size1() const {
return npy_.shape[0];
}
size_t size2() const {
if(npy_.shape.size() == 1)
return 1;
else
return npy_.shape[1];
}
private:
const cnpy::NpyArray& npy_;
};
public:
typedef blaze::CustomMatrix<float, blaze::unaligned,
blaze::unpadded, blaze::rowMajor> BlazeWrapper;
bool has(std::string key) const {
auto it = model_.find(key);
return (it != model_.end());
}
NpzConverter(const std::string& file)
: model_(cnpy::npz_load(file)),
destructed_(false) {
}
~NpzConverter() {
if(!destructed_)
model_.destruct();
}
void Destruct() {
model_.destruct();
destructed_ = true;
}
mblas::Tensor operator[](const std::string& key) const {
BlazeWrapper matrix;
auto it = model_.find(key);
if(it != model_.end()) {
NpyMatrixWrapper np(it->second);
matrix = BlazeWrapper(np.data(), np.size1(), np.size2());
}
else {
if (key.find("gamma") == std::string::npos) {
std::cerr << "Missing " << key << std::endl;
}
}
mblas::Tensor ret;
ret = matrix;
return std::move(ret);
}
mblas::Tensor getFirstOfMany(const std::vector<std::pair<std::string, bool>> keys) const {
BlazeWrapper matrix;
for (auto key : keys) {
auto it = model_.find(key.first);
if(it != model_.end()) {
NpyMatrixWrapper np(it->second);
matrix = BlazeWrapper(np.data(), np.size1(), np.size2());
mblas::Tensor ret;
if (key.second) {
const auto matrix2 = blaze::trans(matrix);
ret = matrix2;
} else {
ret = matrix;
}
return std::move(ret);
}
}
std::cerr << "Matrix not found: " << keys[0].first << "\n";
mblas::Tensor ret;
return std::move(ret);
}
mblas::Tensor operator()(const std::string& key,
bool transpose) const {
BlazeWrapper matrix;
auto it = model_.find(key);
if(it != model_.end()) {
NpyMatrixWrapper np(it->second);
matrix = BlazeWrapper(np.data(), np.size1(), np.size2());
} else {
std::cerr << "Missing " << key << std::endl;
}
mblas::Tensor ret;
if (transpose) {
const auto matrix2 = blaze::trans(matrix);
ret = matrix2;
} else {
ret = matrix;
}
return std::move(ret);
}
private:
cnpy::npz_t model_;
bool destructed_;
};
}
}
|
