-
Notifications
You must be signed in to change notification settings - Fork 1
/
to_v3_converter.py
84 lines (72 loc) · 2.93 KB
/
to_v3_converter.py
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
import argparse
import numpy as np
import chainer
from chainer import serializers
import chainer.functions as F
import chainer.links as L
from models.yolov3_base import YOLOv3_base
from models.yolov3_variants import YOLOv3_update_base
def parse_args():
parser = argparse.ArgumentParser(
description="Converter from darknet to Chainer")
parser.add_argument('--model', default='pre', help='[pre, normal]')
parser.add_argument('--orig', help="File name of darknet's model parameter")
parser.add_argument('--name', help="File name of chainer's model parameter")
parser.add_argument('--n_class', default=80, type=int, help="Number of class")
parser.add_argument('--n_box', default=5, type=int, help="Number of boxes")
parser.add_argument('--darknet', default=0, type=int)
return parser.parse_args()
def to_chainer_converter():
"""model converter from darknet to chainer."""
args = parse_args()
if args.model == 'pre':
Model = YOLOv3_base
elif args.model == 'normal':
Model = YOLOv3_update_base
elif args.model == 'high':
raise("Not Implemented Error: High resolution Yolo")
elif args.model == 'tiny':
raise("Not Implemented Error: Tiny Yolo")
else:
raise("Not Implemented Error")
config = {'n_classes':args.n_class, 'n_boxes':args.n_box}
pretrained_model = {'download': None, 'path': None}
model = Model(config, pretrained_model=pretrained_model)
with open(args.orig, 'rb') as f:
orig_data = np.fromfile(f, dtype='f')[5:] # skip header
i = 1
offset = 0
while True:
try:
if not i in [59, 67, 75]:
bn = getattr(model, 'bn{}'.format(i))
out_ch = bn.gamma.shape[0]
bn.beta.data = orig_data[offset : offset + out_ch]
offset += out_ch
bn.gamma.data = orig_data[offset: offset + out_ch]
offset += out_ch
bn.avg_mean = orig_data[offset: offset + out_ch]
offset += out_ch
bn.avg_var = orig_data[offset: offset + out_ch]
offset += out_ch
conv = getattr(model, 'conv{}'.format(i))
out_ch, in_ch, h, w = conv.W.shape
if i in [59, 67, 75]:
conv.b.data = orig_data[offset: offset+out_ch]
offset += out_ch
print(i, offset, offset + out_ch * in_ch * h * w, orig_data.shape)
conv.W.data = orig_data[offset: offset + out_ch * in_ch * h * w].reshape(out_ch, in_ch, h, w)
offset += out_ch * in_ch * h * w
i += 1
if args.darknet and i == 53:
break
except:
print("Load last convolutional layer", offset, i)
break
save_name = "{}.npz".format(args.name)
chainer.serializers.save_npz(save_name, model)
print("Complete")
def main():
to_chainer_converter()
if __name__ == '__main__':
main()