Skip to content

Commit 7b58065

Browse files
lgarithmlgarithm
committed
add yolo weight converter
1 parent 1bbc71e commit 7b58065

File tree

11 files changed

+232
-12
lines changed

11 files changed

+232
-12
lines changed

‎.gitignore

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010
.DS_store
1111
/.gopath
1212
/.vscode/settings.json
13-
/build
1413
/compile_commands.json
1514
/xtleak.kcg.*
1615
__pycache__
16+
build

‎src/crystalnet-contrib/yolo/conv_layer.cpp

Lines changed: 6 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,8 @@
66
#include <crystalnet/layers/layer.hpp>
77
#include <crystalnet/utility/range.hpp>
88

9+
#include "yolo.hpp"
10+
911
namespace darknet
1012
{
1113
struct conv_op {
@@ -111,15 +113,11 @@ struct conv_layer : s_layer_t {
111113
logf("input: %s", std::to_string(x->shape).c_str());
112114
const auto c = channel_size(x->shape);
113115

114-
const int layer_number = ctx._layers.items.size();
115-
char _name_prefix[32];
116-
sprintf(_name_prefix, "yolov2_%02d_", layer_number);
117-
const std::string name_prefix(_name_prefix);
116+
const std::string prefix = name_prefix(ctx);
118117

119-
const auto weight = ctx.make_parameter(shape_t(size, size, c, filters),
120-
name_prefix + "W");
121-
const auto bias = ctx.make_parameter(shape_t(filters), //
122-
name_prefix + "b");
118+
const auto weight =
119+
ctx.make_parameter(shape_t(size, size, c, filters), prefix + "_W");
120+
const auto bias = ctx.make_parameter(shape_t(filters), prefix + "_b");
123121
return ctx.make_operator(*op, x, weight, bias);
124122
}
125123
};

‎src/crystalnet-contrib/yolo/main.cpp

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ void load_parameters(const model_t *model, const fs::path &model_dir)
2727
const auto filename = model_dir / (name + ".idx");
2828
if (!file_exists(filename)) {
2929
fprintf(stderr, "%s NOT EXIST\n", filename.c_str());
30-
continue;
30+
exit(1);
3131
}
3232
const auto w = _load_idx_file(filename.c_str());
3333
model->ctx.p_ctx.load(name, ref(*w));

‎src/crystalnet-contrib/yolo/region_layer.cpp

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,8 @@
88
#include <crystalnet/layers/layer.hpp>
99
#include <crystalnet/utility/range.hpp>
1010

11+
#include "yolo.hpp"
12+
1113
namespace darknet::yolov2
1214
{
1315
template <typename T> T logistic(T x) { return (T)1. / ((T)1. + std::exp(-x)); }
@@ -117,8 +119,9 @@ struct region_layer : s_layer_t {
117119

118120
s_node_t *operator()(s_model_ctx_t &ctx, s_node_t *x) const override
119121
{
120-
symbol bias = ctx.make_parameter(shape_t(n, 2), "yolov2_anchors");
121-
return ctx.make_operator(*op, x, bias);
122+
const std::string prefix = name_prefix(ctx);
123+
symbol anchors = ctx.make_parameter(shape_t(n, 2), prefix + "_anchors");
124+
return ctx.make_operator(*op, x, anchors);
122125
}
123126
};
124127
} // namespace darknet::yolov2

‎src/crystalnet-contrib/yolo/yolo.hpp

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
#pragma once
2+
#include <string>
3+
4+
#include <crystalnet-ext.h>
5+
#include <crystalnet/symbol/model.hpp>
6+
7+
inline std::string name_prefix(const s_model_ctx_t &ctx)
8+
{
9+
const int layer_number = ctx._layers.items.size();
10+
char name_prefix[32];
11+
sprintf(name_prefix, "yolov2_%02d", layer_number);
12+
return name_prefix;
13+
}

‎utils/light/CMakeLists.txt

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
cmake_minimum_required(VERSION 3.9)
2+
project(light)
3+
4+
set(CMAKE_CXX_STANDARD 17)
5+
6+
set(CRYSTALNET_PATH $ENV{HOME}/local/crystalnet) # make it configurable
7+
include_directories(${CRYSTALNET_PATH}/include)
8+
link_directories(${CRYSTALNET_PATH}/lib)
9+
10+
set(DARKNET_PATH $ENV{HOME}/local/darknet) # make it configurable
11+
include_directories(${DARKNET_PATH}/include)
12+
link_directories(${DARKNET_PATH}/lib)
13+
14+
file(GLOB srcs src/*.cpp)
15+
add_executable(light ${srcs})
16+
target_link_libraries(light crystalnet darknet)

‎utils/light/Makefile

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
ROOT = $(shell pwd)
2+
BUILD_DIR = $(ROOT)/build/$(shell uname)
3+
4+
ifeq ($(shell uname), Darwin)
5+
NPROC = $(shell sysctl -n hw.ncpu)
6+
else
7+
NPROC = $(shell nproc)
8+
endif
9+
10+
default: light
11+
12+
CMAKE_FLAGS = \
13+
-DCMAKE_EXPORT_COMPILE_COMMANDS=1 \
14+
-DCMAKE_BUILD_TYPE=Release \
15+
16+
cmake_targets:
17+
mkdir -p $(BUILD_DIR)
18+
cd $(BUILD_DIR); cmake $(CMAKE_FLAGS) $(ROOT)
19+
20+
light: cmake_targets
21+
make -C $(BUILD_DIR) -j $(NPROC)

‎utils/light/README.md

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
light
2+
======
3+
4+
Light is small tool to convert darknet parameters to idx format.

‎utils/light/src/light.cpp

Lines changed: 122 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,122 @@
1+
#include <cstdio>
2+
#include <cstdlib>
3+
#include <cstring>
4+
#include <experimental/filesystem>
5+
#include <getopt.h>
6+
7+
namespace fs = std::experimental::filesystem;
8+
9+
#include <crystalnet.h>
10+
11+
extern "C" {
12+
#include <darknet.h>
13+
extern char *get_layer_string(LAYER_TYPE a);
14+
}
15+
16+
#include "options.hpp"
17+
18+
std::string layer_name(LAYER_TYPE lt) { return get_layer_string(lt); }
19+
20+
struct saver_t {
21+
const fs::path path;
22+
saver_t(const fs::path &path) : path(path) {}
23+
24+
void operator()(const std::string &name, const tensor_ref_t *r) const
25+
{
26+
const auto filename = path / (name + ".idx");
27+
save_tensor(filename.c_str(), r);
28+
printf("[i] %s saved to %s\n", name.c_str(), filename.c_str());
29+
}
30+
};
31+
32+
struct converter_t {
33+
const std::string name_prefix;
34+
const saver_t save;
35+
36+
converter_t(const fs::path &path, const std::string &name_prefix)
37+
: save(path), name_prefix(name_prefix)
38+
{
39+
}
40+
41+
void save_conv_layer(const std::string &name, const layer &l) const
42+
{
43+
const shape_t *bias_shape = new_shape(1, l.n);
44+
const shape_t *weight_shape =
45+
new_shape(4, l.size, l.size, l.c, l.n); // TODO: double check
46+
if (l.nweights != shape_dim(weight_shape)) {
47+
fprintf(stderr, "invalid layer\n");
48+
exit(1);
49+
}
50+
const tensor_t *bias = new_tensor(bias_shape, dtypes.f32);
51+
const tensor_t *weight = new_tensor(weight_shape, dtypes.f32);
52+
std::memcpy(tensor_data_ptr(tensor_ref(bias)), l.biases,
53+
shape_dim(bias_shape) * sizeof(float));
54+
std::memcpy(tensor_data_ptr(tensor_ref(weight)), l.weights,
55+
shape_dim(weight_shape) * sizeof(float));
56+
save(name + "_b", tensor_ref(bias));
57+
save(name + "_W", tensor_ref(weight));
58+
del_tensor(bias);
59+
del_tensor(weight);
60+
del_shape(bias_shape);
61+
del_shape(weight_shape);
62+
}
63+
64+
void save_region_layer(const std::string &name, const layer &l) const
65+
{
66+
const shape_t *shape = new_shape(2, l.n, 2);
67+
const tensor_t *anchors = new_tensor(shape, dtypes.f32);
68+
save(name + "_anchors", tensor_ref(anchors));
69+
del_tensor(anchors);
70+
del_shape(shape);
71+
}
72+
73+
std::string lpad(const std::string s, int width, char ch) const
74+
{
75+
return std::string(std::max<int>(width - s.size(), 0), ch) + s;
76+
}
77+
78+
void save_parameters(const network *net) const
79+
{
80+
printf("saving weights ...\n");
81+
for (int i = 0; i < net->n; ++i) {
82+
const auto l = net->layers[i];
83+
const std::string prefix =
84+
name_prefix + "_" + lpad(std::to_string(i), 2, '0');
85+
switch (l.type) {
86+
case CONVOLUTIONAL:
87+
save_conv_layer(prefix, l);
88+
continue;
89+
case REGION:
90+
save_region_layer(prefix, l);
91+
continue;
92+
case MAXPOOL:
93+
case REORG:
94+
case ROUTE:
95+
printf("[d] %s layer has no weights\n",
96+
layer_name(l.type).c_str());
97+
continue;
98+
default:
99+
// TODO: save other layers
100+
printf("TODO: layer %-4d: %s\n", i, layer_name(l.type).c_str());
101+
}
102+
}
103+
printf("done.\n");
104+
}
105+
106+
void operator()(const fs::path &model_cfg,
107+
const fs::path &weights_file) const
108+
{
109+
network *net = load_network((char *)model_cfg.c_str(),
110+
(char *)weights_file.c_str(), 0);
111+
printf("%d layers\n", net->n);
112+
save_parameters(net);
113+
}
114+
};
115+
116+
int main(int argc, char *argv[])
117+
{
118+
const auto opts = parse_flags(argc, argv);
119+
converter_t convert(opts.model_dir, "yolov2");
120+
convert(opts.model_cfg, opts.weights_file);
121+
return 0;
122+
}

‎utils/light/src/options.cpp

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
#include "options.hpp"
2+
3+
#include <getopt.h>
4+
5+
namespace fs = std::experimental::filesystem;
6+
7+
options_t::options_t()
8+
{
9+
const fs::path home(std::getenv("HOME"));
10+
const auto darknet_path = home / "code/mirrors/github.com/pjreddie/darknet";
11+
model_dir = home / "var/models/yolo";
12+
model_cfg = darknet_path / "cfg/yolov2.cfg";
13+
weights_file = model_dir / "yolov2.weights";
14+
}
15+
16+
options_t parse_flags(int argc, char *argv[])
17+
{
18+
options_t option;
19+
for (int c; (c = getopt(argc, argv, "c:")) != -1;) {
20+
switch (c) {
21+
case 'c':
22+
option.model_cfg = optarg;
23+
break;
24+
default:
25+
fprintf(stderr, "unknown options: %c", c);
26+
break;
27+
}
28+
}
29+
return option;
30+
}

‎utils/light/src/options.hpp

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
#pragma once
2+
3+
#include <experimental/filesystem>
4+
5+
struct options_t {
6+
std::experimental::filesystem::path model_cfg;
7+
std::experimental::filesystem::path model_dir;
8+
std::experimental::filesystem::path weights_file;
9+
10+
options_t();
11+
};
12+
13+
options_t parse_flags(int /* argc */, char *argv[]);

0 commit comments

Comments
 (0)