-
Notifications
You must be signed in to change notification settings - Fork 83
/
train_model.py
171 lines (134 loc) · 8.42 KB
/
train_model.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
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
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
#!/usr/bin/env python
#####################################################################################################################################################################
# xView2 #
# Copyright 2019 Carnegie Mellon University. #
# NO WARRANTY. THIS CARNEGIE MELLON UNIVERSITY AND SOFTWARE ENGINEERING INSTITUTE MATERIAL IS FURNISHED ON AN "AS-IS" BASIS. CARNEGIE MELLON UNIVERSITY MAKES NO #
# WARRANTIES OF ANY KIND, EITHER EXPRESSED OR IMPLIED, AS TO ANY MATTER INCLUDING, BUT NOT LIMITED TO, WARRANTY OF FITNESS FOR PURPOSE OR MERCHANTABILITY, #
# EXCLUSIVITY, OR RESULTS OBTAINED FROM USE OF THE MATERIAL. CARNEGIE MELLON UNIVERSITY DOES NOT MAKE ANY WARRANTY OF ANY KIND WITH RESPECT TO FREEDOM FROM PATENT, #
# TRADEMARK, OR COPYRIGHT INFRINGEMENT. #
# Released under a MIT (SEI)-style license, please see LICENSE.md or contact [email protected] for full terms. #
# [DISTRIBUTION STATEMENT A] This material has been approved for public release and unlimited distribution. Please see Copyright notice for non-US Government use #
# and distribution. #
# This Software includes and/or makes use of the following Third-Party Software subject to its own license: #
# 1. SpaceNet (https://github.com/motokimura/spacenet_building_detection/blob/master/LICENSE) Copyright 2017 Motoki Kimura. #
# DM19-0988 #
#####################################################################################################################################################################
from __future__ import print_function
import argparse
import numpy as np
import chainer
import chainer.functions as F
import chainer.links as L
from chainer import training
from chainer.training import extensions
from unet import UNet
from dataset import LabeledImageDataset
from tensorboardX import SummaryWriter
import os
def train_model():
parser = argparse.ArgumentParser()
parser.add_argument('dataset', help='Path to directory containing train.txt, val.txt, and mean.npy')
parser.add_argument('images', help='Root directory of input images')
parser.add_argument('labels', help='Root directory of label images')
parser.add_argument('--batchsize', '-b', type=int, default=16,
help='Number of images in each mini-batch')
parser.add_argument('--test-batchsize', '-B', type=int, default=4,
help='Number of images in each test mini-batch')
parser.add_argument('--epoch', '-e', type=int, default=50,
help='Number of sweeps over the dataset to train')
parser.add_argument('--frequency', '-f', type=int, default=1,
help='Frequency of taking a snapshot')
parser.add_argument('--gpu', '-g', type=int, default=0,
help='GPU ID (negative value indicates CPU)')
parser.add_argument('--out', '-o', default='logs',
help='Directory to output the result under "models" directory')
parser.add_argument('--resume', '-r', default='',
help='Resume the training from snapshot')
parser.add_argument('--noplot', dest='plot', action='store_false',
help='Disable PlotReport extension')
parser.add_argument('--tcrop', type=int, default=400,
help='Crop size for train-set images')
parser.add_argument('--vcrop', type=int, default=480,
help='Crop size for validation-set images')
args = parser.parse_args()
assert (args.tcrop % 16 == 0) and (args.vcrop % 16 == 0), "tcrop and vcrop must be divisible by 16."
if args.gpu < 0:
from tboard_logger_cpu import TensorboardLogger
else:
from tboard_logger import TensorboardLogger
print('GPU: {}'.format(args.gpu))
print('# Minibatch-size: {}'.format(args.batchsize))
print('# Crop-size: {}'.format(args.tcrop))
print('# epoch: {}'.format(args.epoch))
print('')
this_dir = os.path.dirname(os.path.abspath(__file__))
models_dir = os.path.normpath(os.path.join(this_dir, "../../models"))
log_dir = os.path.join(models_dir, args.out)
writer = SummaryWriter(log_dir=log_dir)
# Set up a neural network to train
# Classifier reports softmax cross entropy loss and accuracy at every
# iteration, which will be used by the PrintReport extension below.
model = UNet()
if args.gpu >= 0:
# Make a specified GPU current
chainer.cuda.get_device_from_id(args.gpu).use()
model.to_gpu() # Copy the model to the GPU
# Setup an optimizer
optimizer = chainer.optimizers.Adam()
optimizer.setup(model)
# Load mean image
mean = np.load(os.path.join(args.dataset, "mean.npy"))
# Load the MNIST dataset
train = LabeledImageDataset(os.path.join(args.dataset, "train.txt"), args.images, args.labels,
mean=mean, crop_size=args.tcrop, test=False, distort=False)
test = LabeledImageDataset (os.path.join(args.dataset, "val.txt"), args.images, args.labels,
mean=mean, crop_size=args.vcrop, test=True, distort=False)
train_iter = chainer.iterators.SerialIterator(train, args.batchsize)
test_iter = chainer.iterators.SerialIterator(test, args.test_batchsize, repeat=False, shuffle=False)
# Set up a trainer
updater = training.StandardUpdater(
train_iter, optimizer, device=args.gpu)
trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=log_dir)
# Evaluate the model with the test dataset for each epoch
trainer.extend(extensions.Evaluator(test_iter, model, device=args.gpu))
# Dump a computational graph from 'loss' variable at the first iteration
# The "main" refers to the target link of the "main" optimizer.
trainer.extend(extensions.dump_graph('main/loss'))
# Take a snapshot for each specified epoch
frequency = args.epoch if args.frequency == -1 else max(1, args.frequency)
trainer.extend(extensions.snapshot(), trigger=(frequency, 'epoch'))
# Save trained model for each specific epoch
trainer.extend(extensions.snapshot_object(
model, 'model_iter_{.updater.iteration}'), trigger=(frequency, 'epoch'))
# Write a log of evaluation statistics for each epoch
trainer.extend(extensions.LogReport())
# Save two plot images to the result dir
if args.plot and extensions.PlotReport.available():
trainer.extend(
extensions.PlotReport(['main/loss', 'validation/main/loss'],
'epoch', file_name='loss.png'))
trainer.extend(
extensions.PlotReport(
['main/accuracy', 'validation/main/accuracy'],
'epoch', file_name='accuracy.png'))
# Print selected entries of the log to stdout
# Here "main" refers to the target link of the "main" optimizer again, and
# "validation" refers to the default name of the Evaluator extension.
# Entries other than 'epoch' are reported by the Classifier link, called by
# either the updater or the evaluator.
trainer.extend(extensions.PrintReport(
['epoch', 'main/loss', 'validation/main/loss',
'main/accuracy', 'validation/main/accuracy', 'elapsed_time']))
# Print a progress bar to stdout
trainer.extend(extensions.ProgressBar())
# Write training log to TensorBoard log file
trainer.extend(TensorboardLogger(writer,
['main/loss', 'validation/main/loss',
'main/accuracy', 'validation/main/accuracy']))
if args.resume:
# Resume from a snapshot
chainer.serializers.load_npz(args.resume, trainer)
# Run the training
trainer.run()
if __name__ == '__main__':
train_model()