-
Notifications
You must be signed in to change notification settings - Fork 1
/
Copy pathtrain.py
60 lines (48 loc) · 2.08 KB
/
train.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
#!/usr/env/bin python3
# -*- coding: utf-8 -*-
import argparse
import numpy as np
import sys
import subprocess
import os
import yaml
subprocess.call(['sh', 'setup.sh'])
import chainer
from chainer import cuda, optimizers, serializers
from chainer import training
from chainer.datasets import TransformDataset
from chainercv.links.model.ssd import random_distort
from config_utils import *
from datasets.transform import Transform as Transform_v2
from datasets.v3_transform import Transform as Transform_v3
chainer.cuda.set_max_workspace_size(1024 * 1024 * 1024)
os.environ["CHAINER_TYPE_CHECK"] = "0"
# chainer.global_config.debug = True
from collections import OrderedDict
yaml.add_constructor(yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG,
lambda loader, node: OrderedDict(loader.construct_pairs(node)))
def train_yolov2():
"""Training yolov2."""
config = parse_args()
model = get_model(config["model"])
devices = parse_devices(config['gpus'], config['updater']['name'])
train_data, test_data = load_dataset(config["dataset"])
Transform = Transform_v2 if parse_dict(config, 'version', '2') == '2' else Transform_v3
train_data = TransformDataset(
train_data, Transform(0.5, dim=model.dim, max_target=30,
anchors=model.anchors, batchsize=config['iterator']['train_batchsize']))
train_iter, test_iter = create_iterator(train_data, test_data,
config['iterator'], devices,
config['updater']['name'])
optimizer = create_optimizer(config['optimizer'], model)
updater = create_updater(train_iter, optimizer, config['updater'], devices)
trainer = training.Trainer(updater, config['end_trigger'], out=config['results'])
trainer = create_extension(trainer, test_iter, model,
config['extension'], devices=devices)
trainer.run()
chainer.serializers.save_npz(os.path.join(config['results'], 'model.npz'),
model)
def main():
train_yolov2()
if __name__ == '__main__':
main()