def infer():
# check if set use_gpu=True in paddlepaddle cpu version
check_gpu(cfg.use_gpu)
if not os.path.exists('output'):
os.mkdir('output')
model = YOLOv3(is_train=False)
model.build_model()
outputs = model.get_pred()
input_size = cfg.input_size
place = fluid.CUDAPlace(0) if cfg.use_gpu else fluid.CPUPlace()
exe = fluid.Executor(place)
# yapf: disable
if cfg.weights:
def if_exist(var):
return os.path.exists(os.path.join(cfg.weights, var.name))
fluid.io.load_vars(exe, cfg.weights, predicate=if_exist)
# yapf: enable
# you can save inference model by following code
# fluid.io.save_inference_model("./output/yolov3",
# feeded_var_names=['image', 'im_shape'],
# target_vars=outputs,
# executor=exe)
feeder = fluid.DataFeeder(place=place, feed_list=model.feeds())
fetch_list = [outputs]
image_names = []
if cfg.image_name is not None:
image_names.append(cfg.image_name)
else:
for image_name in os.listdir(cfg.image_path):
if image_name.split('.')[-1] in ['jpg', 'png']:
image_names.append(image_name)
for image_name in image_names:
infer_reader = reader.infer(input_size,
os.path.join(cfg.image_path, image_name))
label_names, _ = reader.get_label_infos()
data = next(infer_reader())
im_shape = data[0][2]
outputs = exe.run(fetch_list=[v.name for v in fetch_list],
feed=feeder.feed(data),
return_numpy=False,
use_program_cache=True)
bboxes = np.array(outputs[0])
if bboxes.shape[1] != 6:
print("No object found in {}".format(image_name))
continue
labels = bboxes[:, 0].astype('int32')
scores = bboxes[:, 1].astype('float32')
boxes = bboxes[:, 2:].astype('float32')
path = os.path.join(cfg.image_path, image_name)
box_utils.draw_boxes_on_image(path, boxes, scores, labels, label_names,
cfg.draw_thresh)
def infer():
# check if set use_gpu=True in paddlepaddle cpu version
check_gpu(cfg.use_gpu)
if not os.path.exists('output'):
os.mkdir('output')
place = fluid.CUDAPlace(0) if cfg.use_gpu else fluid.CPUPlace()
with fluid.dygraph.guard(place):
model = YOLOv3(3, is_train=False)
input_size = cfg.input_size
# yapf: disable
if cfg.weights:
restore, _ = fluid.load_dygraph(cfg.weights)
model.set_dict(restore)
# yapf: enable
# you can save inference model by following code
# fluid.io.save_inference_model("./output/yolov3",
# feeded_var_names=['image', 'im_shape'],
# target_vars=outputs,
# executor=exe)
image_names = []
if cfg.image_name is not None:
image_names.append(cfg.image_name)
else:
for image_name in os.listdir(cfg.image_path):
if image_name.split('.')[-1] in ['jpg', 'png']:
image_names.append(image_name)
for image_name in image_names:
infer_reader = reader.infer(
input_size, os.path.join(cfg.image_path, image_name))
label_names, _ = reader.get_label_infos()
data = next(infer_reader())
img_data = np.array([x[0] for x in data]).astype('float32')
img = to_variable(img_data)
im_shape_data = np.array([x[2] for x in data]).astype('int32')
im_shape = to_variable(im_shape_data)
outputs = model(img, None, None, None, None, im_shape)
bboxes = outputs.numpy()
if bboxes.shape[1] != 6:
print("No object found in {}".format(image_name))
continue
labels = bboxes[:, 0].astype('int32')
scores = bboxes[:, 1].astype('float32')
boxes = bboxes[:, 2:].astype('float32')
path = os.path.join(cfg.image_path, image_name)
box_utils.draw_boxes_on_image(path, boxes, scores, labels,
label_names, cfg.draw_thresh)
def infer():
if not os.path.exists('output'):
os.mkdir('output')
model = YOLOv3(is_train=False)
model.build_model()
outputs = model.get_pred()
input_size = cfg.input_size
#place = fluid.CUDAPlace(0) if cfg.use_gpu else fluid.CPUPlace()
place = fluid.CPUPlace()
exe = fluid.Executor(place)
# yapf: disable
if cfg.weights:
print('cfg.weights = ',cfg.weights)
def if_exist(var):
return os.path.exists(os.path.join(cfg.weights, var.name))
fluid.io.load_vars(exe, cfg.weights, predicate=if_exist)
# yapf: enable
feeder = fluid.DataFeeder(place=place, feed_list=model.feeds())
fetch_list = [outputs]
image_names = []
if cfg.image_name is not None:
image_names.append(cfg.image_name)
print('image_names = ', image_names)
else:
for image_name in os.listdir(cfg.image_path):
if image_name.split('.')[-1] in ['jpg', 'png']:
image_names.append(image_name)
for image_name in image_names:
infer_reader = reader.infer(input_size,
os.path.join(cfg.image_path, image_name))
label_names, _ = reader.get_label_infos()
data = next(infer_reader())
im_shape = data[0][2]
outputs = exe.run(fetch_list=[v.name for v in fetch_list],
feed=feeder.feed(data),
return_numpy=False)
bboxes = np.array(outputs[0])
if bboxes.shape[1] != 6:
print("No object found in {}".format(image_name))
continue
labels = bboxes[:, 0].astype('int32')
scores = bboxes[:, 1].astype('float32')
boxes = bboxes[:, 2:].astype('float32')
print(labels)
print(scores)
print(len(boxes))
path = os.path.join(cfg.image_path, image_name)
box_utils.draw_boxes_on_image(path, boxes, scores, labels, label_names,
cfg.draw_thresh)
def load_model(self):
"""
Load the model using settings.CONFIG_FILE and settings.MODEL_CHECKPOINT
"""
with open(settings.CONFIG_FILE, 'r') as f:
cfg = yaml.load(f)
self.imgsize = cfg['TEST']['IMGSIZE']
self.model = YOLOv3(cfg['MODEL'])
self.confthre = cfg['TEST']['CONFTHRE']
self.nmsthre = cfg['TEST']['NMSTHRE']
if use_cuda():
print("Using cuda")
self.model = self.model.cuda()
print("Loading checkpoint {}".format(settings.MODEL_CHECKPOINT))
state = torch.load(settings.MODEL_CHECKPOINT)
if 'model_state_dict' in state.keys():
self.model.load_state_dict(state['model_state_dict'])
else:
self.model.load_state_dict(state)
self.model.eval()
def train():
if cfg.debug or args.enable_ce:
fluid.default_startup_program().random_seed = 1000
fluid.default_main_program().random_seed = 1000
random.seed(0)
np.random.seed(0)
if not os.path.exists(cfg.model_save_dir):
os.makedirs(cfg.model_save_dir)
model = YOLOv3()
model.build_model()
input_size = cfg.input_size
loss = model.loss()
loss.persistable = True
devices_num = get_device_num()
print("Found {} CUDA devices.".format(devices_num))
learning_rate = cfg.learning_rate
boundaries = cfg.lr_steps
gamma = cfg.lr_gamma
step_num = len(cfg.lr_steps)
values = [learning_rate * (gamma**i) for i in range(step_num + 1)]
optimizer = fluid.optimizer.Momentum(
learning_rate=exponential_with_warmup_decay(
learning_rate=learning_rate,
boundaries=boundaries,
values=values,
warmup_iter=cfg.warm_up_iter,
warmup_factor=cfg.warm_up_factor),
regularization=fluid.regularizer.L2Decay(cfg.weight_decay),
momentum=cfg.momentum)
optimizer.minimize(loss)
gpu_id = int(os.environ.get('FLAGS_selected_gpus', 0))
place = fluid.CUDAPlace(gpu_id) if cfg.use_gpu else fluid.CPUPlace()
exe = fluid.Executor(place)
exe.run(fluid.default_startup_program())
if cfg.pretrain:
if not os.path.exists(cfg.pretrain):
print("Pretrain weights not found: {}".format(cfg.pretrain))
def if_exist(var):
return os.path.exists(os.path.join(cfg.pretrain, var.name)) \
and var.name.find('yolo_output') < 0
fluid.io.load_vars(exe, cfg.pretrain, predicate=if_exist)
build_strategy = fluid.BuildStrategy()
build_strategy.memory_optimize = False #gc and memory optimize may conflict
syncbn = cfg.syncbn
if (syncbn and devices_num <= 1) or num_trainers > 1:
print("Disable syncbn in single device")
syncbn = False
build_strategy.sync_batch_norm = syncbn
exec_strategy = fluid.ExecutionStrategy()
if cfg.use_gpu and num_trainers > 1:
dist_utils.prepare_for_multi_process(exe, build_strategy,
fluid.default_main_program())
exec_strategy.num_threads = 1
compile_program = fluid.compiler.CompiledProgram(fluid.default_main_program(
)).with_data_parallel(
loss_name=loss.name,
build_strategy=build_strategy,
exec_strategy=exec_strategy)
random_sizes = [cfg.input_size]
if cfg.random_shape:
random_sizes = [32 * i for i in range(10, 20)]
total_iter = cfg.max_iter - cfg.start_iter
mixup_iter = total_iter - cfg.no_mixup_iter
shuffle = True
if args.enable_ce:
shuffle = False
shuffle_seed = None
# NOTE: yolov3 is a special model, if num_trainers > 1, each process
# trian the completed dataset.
# if num_trainers > 1: shuffle_seed = 1
train_reader = reader.train(
input_size,
batch_size=cfg.batch_size,
shuffle=shuffle,
shuffle_seed=shuffle_seed,
total_iter=total_iter * devices_num,
mixup_iter=mixup_iter * devices_num,
random_sizes=random_sizes,
use_multiprocess_reader=cfg.use_multiprocess_reader)
py_reader = model.py_reader
py_reader.decorate_paddle_reader(train_reader)
def save_model(postfix):
model_path = os.path.join(cfg.model_save_dir, postfix)
if os.path.isdir(model_path):
shutil.rmtree(model_path)
fluid.io.save_persistables(exe, model_path)
fetch_list = [loss]
py_reader.start()
smoothed_loss = SmoothedValue()
try:
start_time = time.time()
prev_start_time = start_time
snapshot_loss = 0
snapshot_time = 0
for iter_id in range(cfg.start_iter, cfg.max_iter):
prev_start_time = start_time
start_time = time.time()
losses = exe.run(compile_program,
fetch_list=[v.name for v in fetch_list])
smoothed_loss.add_value(np.mean(np.array(losses[0])))
snapshot_loss += np.mean(np.array(losses[0]))
snapshot_time += start_time - prev_start_time
lr = np.array(fluid.global_scope().find_var('learning_rate')
.get_tensor())
print("Iter {:d}, lr {:.6f}, loss {:.6f}, time {:.5f}".format(
iter_id, lr[0],
smoothed_loss.get_mean_value(), start_time - prev_start_time))
sys.stdout.flush()
if (iter_id + 1) % cfg.snapshot_iter == 0:
save_model("model_iter{}".format(iter_id))
print("Snapshot {} saved, average loss: {}, \
average time: {}".format(
iter_id + 1, snapshot_loss / float(cfg.snapshot_iter),
snapshot_time / float(cfg.snapshot_iter)))
if args.enable_ce and iter_id == cfg.max_iter - 1:
if devices_num == 1:
print("kpis\ttrain_cost_1card\t%f" %
(snapshot_loss / float(cfg.snapshot_iter)))
print("kpis\ttrain_duration_1card\t%f" %
(snapshot_time / float(cfg.snapshot_iter)))
else:
print("kpis\ttrain_cost_8card\t%f" %
(snapshot_loss / float(cfg.snapshot_iter)))
print("kpis\ttrain_duration_8card\t%f" %
(snapshot_time / float(cfg.snapshot_iter)))
snapshot_loss = 0
snapshot_time = 0
except fluid.core.EOFException:
py_reader.reset()
save_model('model_final')
def eval():
# check if set use_gpu=True in paddlepaddle cpu version
check_gpu(cfg.use_gpu)
if '2014' in cfg.dataset:
test_list = 'annotations/instances_val2014.json'
elif '2017' in cfg.dataset:
test_list = 'annotations/instances_val2017.json'
if cfg.debug:
if not os.path.exists('output'):
os.mkdir('output')
place = fluid.CUDAPlace(0) if cfg.use_gpu else fluid.CPUPlace()
with fluid.dygraph.guard(place):
model = YOLOv3(3,is_train=False)
# yapf: disable
if cfg.weights:
restore, _ = fluid.load_dygraph(cfg.weights)
model.set_dict(restore)
model.eval()
input_size = cfg.input_size
# batch_size for test must be 1
test_reader = reader.test(input_size, 1)
label_names, label_ids = reader.get_label_infos()
if cfg.debug:
print("Load in labels {} with ids {}".format(label_names, label_ids))
def get_pred_result(boxes, scores, labels, im_id):
result = []
for box, score, label in zip(boxes, scores, labels):
x1, y1, x2, y2 = box
w = x2 - x1 + 1
h = y2 - y1 + 1
bbox = [x1, y1, w, h]
res = {
'image_id': int(im_id),
'category_id': label_ids[int(label)],
'bbox': list(map(float, bbox)),
'score': float(score)
}
result.append(res)
return result
dts_res = []
total_time = 0
for iter_id, data in enumerate(test_reader()):
start_time = time.time()
img_data = np.array([x[0] for x in data]).astype('float32')
img = to_variable(img_data)
im_id_data = np.array([x[1] for x in data]).astype('int32')
im_id = to_variable(im_id_data)
im_shape_data = np.array([x[2] for x in data]).astype('int32')
im_shape = to_variable(im_shape_data)
batch_outputs = model(img, None, None, None, im_id, im_shape)
nmsed_boxes = batch_outputs.numpy()
if nmsed_boxes.shape[1] != 6:
continue
im_id = data[0][1]
nmsed_box=nmsed_boxes
labels = nmsed_box[:, 0]
scores = nmsed_box[:, 1]
boxes = nmsed_box[:, 2:6]
dts_res += get_pred_result(boxes, scores, labels, im_id)
end_time = time.time()
print("batch id: {}, time: {}".format(iter_id, end_time - start_time))
total_time += end_time - start_time
with open("yolov3_result.json", 'w') as outfile:
json.dump(dts_res, outfile)
print("start evaluate detection result with coco api")
coco = COCO(os.path.join(cfg.data_dir, test_list))
cocoDt = coco.loadRes("yolov3_result.json")
cocoEval = COCOeval(coco, cocoDt, 'bbox')
cocoEval.evaluate()
cocoEval.accumulate()
cocoEval.summarize()
print("evaluate done.")
print("Time per batch: {}".format(total_time / iter_id))
def infer():
if not os.path.exists('output_frame'):
os.mkdir('output_frame')
model = YOLOv3(is_train=False)
model.build_model()
outputs = model.get_pred()
input_size = cfg.input_size
place = fluid.CUDAPlace(0) if cfg.use_gpu else fluid.CPUPlace()
exe = fluid.Executor(place)
# yapf: disable
if cfg.weights:
def if_exist(var):
return os.path.exists(os.path.join(cfg.weights, var.name))
fluid.io.load_vars(exe, cfg.weights, predicate=if_exist)
# yapf: enable
# you can save inference model by following code
# fluid.io.save_inference_model("./output/yolov3",
# feeded_var_names=['image', 'im_shape'],
# target_vars=outputs,
# executor=exe)
feeder = fluid.DataFeeder(place=place, feed_list=model.feeds())
fetch_list = [outputs]
video_path = "" #source
video_dir = "" #output path
video_writer = FFmpegWriter(video_dir)
cap = cv2.VideoCapture(video_path)
i = 0
while cap.isOpened():
if i % 100 == 0:
print('{}% '.format(i / 1020 * 100) + '%d' % i)
label_names, _ = reader.get_label_infos()
ok, image = cap.read()
if ok:
infer_reader = reader.infer_m(input_size, image)
data = next(infer_reader())
im_shape = data[0][2]
outputs = exe.run(fetch_list=[v.name for v in fetch_list],
feed=feeder.feed(data),
return_numpy=False)
bboxes = np.array(outputs[0])
if bboxes.shape[1] != 6:
continue
labels = bboxes[:, 0].astype('int32')
scores = bboxes[:, 1].astype('float32')
boxes = bboxes[:, 2:].astype('float32')
frame = box_utils.save_as_video(image, boxes, scores, labels,
label_names, cfg.draw_thresh)
frame = frame[:, :, ::-1]
frame = cv2.resize(frame, (492, 369))
video_writer.writeFrame(frame)
else:
print("over")
break
i += 1
video_writer.close()
def train():
# check if set use_gpu=True in paddlepaddle cpu version
check_gpu(cfg.use_gpu)
devices_num = get_device_num() if cfg.use_gpu else 1
print("Found {} CUDA/CPU devices.".format(devices_num))
if cfg.debug or args.enable_ce:
fluid.default_startup_program().random_seed = 1000
fluid.default_main_program().random_seed = 1000
random.seed(0)
np.random.seed(0)
if not os.path.exists(cfg.model_save_dir):
os.makedirs(cfg.model_save_dir)
gpu_id = int(os.environ.get('FLAGS_selected_gpus', 0))
place = fluid.CUDAPlace(fluid.dygraph.parallel.Env().dev_id) if cfg.use_data_parallel else fluid.CUDAPlace(0)
with fluid.dygraph.guard(place):
if args.use_data_parallel:
strategy = fluid.dygraph.parallel.prepare_context()
model = YOLOv3(3, is_train=True)
if cfg.pretrain:
restore, _ = fluid.load_dygraph(cfg.pretrain)
model.block.set_dict(restore)
if cfg.finetune:
restore, _ = fluid.load_dygraph(cfg.finetune)
model.set_dict(restore, use_structured_name=True)
if args.use_data_parallel:
model = fluid.dygraph.parallel.DataParallel(model, strategy)
boundaries = cfg.lr_steps
gamma = cfg.lr_gamma
step_num = len(cfg.lr_steps)
learning_rate = cfg.learning_rate
values = [learning_rate * (gamma ** i) for i in range(step_num + 1)]
lr = fluid.dygraph.PiecewiseDecay(
boundaries=boundaries,
values=values,
begin=args.start_iter)
lr = fluid.layers.linear_lr_warmup(
learning_rate=lr,
warmup_steps=cfg.warm_up_iter,
start_lr=0.0,
end_lr=cfg.learning_rate,
)
optimizer = fluid.optimizer.Momentum(
learning_rate=lr,
regularization=fluid.regularizer.L2Decay(cfg.weight_decay),
momentum=cfg.momentum,
parameter_list=model.parameters()
)
start_time = time.time()
snapshot_loss = 0
snapshot_time = 0
total_sample = 0
input_size = cfg.input_size
shuffle = True
shuffle_seed = None
total_iter = cfg.max_iter - cfg.start_iter
mixup_iter = total_iter - cfg.no_mixup_iter
random_sizes = [cfg.input_size]
if cfg.random_shape:
random_sizes = [32 * i for i in range(10,20)]
train_reader = reader.train(
input_size,
batch_size=cfg.batch_size,
shuffle=shuffle,
shuffle_seed=shuffle_seed,
total_iter=total_iter * devices_num,
mixup_iter=mixup_iter * devices_num,
random_sizes=random_sizes,
use_multiprocess_reader=cfg.use_multiprocess_reader,
num_workers=cfg.worker_num)
if args.use_data_parallel:
train_reader = fluid.contrib.reader.distributed_batch_reader(train_reader)
smoothed_loss = SmoothedValue()
for iter_id, data in enumerate(train_reader()):
prev_start_time = start_time
start_time = time.time()
img = np.array([x[0] for x in data]).astype('float32')
img = to_variable(img)
gt_box = np.array([x[1] for x in data]).astype('float32')
gt_box = to_variable(gt_box)
gt_label = np.array([x[2] for x in data]).astype('int32')
gt_label = to_variable(gt_label)
gt_score = np.array([x[3] for x in data]).astype('float32')
gt_score = to_variable(gt_score)
loss = model(img, gt_box, gt_label, gt_score, None, None)
smoothed_loss.add_value(np.mean(loss.numpy()))
snapshot_loss += loss.numpy()
snapshot_time += start_time - prev_start_time
total_sample += 1
print("Iter {:d}, loss {:.6f}, time {:.5f}".format(
iter_id,
smoothed_loss.get_mean_value(),
start_time-prev_start_time))
if args.use_data_parallel:
loss = model.scale_loss(loss)
loss.backward()
model.apply_collective_grads()
loss.backward()
optimizer.minimize(loss)
model.clear_gradients()
save_parameters = (not args.use_data_parallel) or (
args.use_data_parallel and
fluid.dygraph.parallel.Env().local_rank == 0)
if save_parameters and iter_id > 1 and iter_id % cfg.snapshot_iter == 0:
fluid.save_dygraph(model.state_dict(), args.model_save_dir + "/yolov3_{}".format(iter_id))
def eval():
if '2014' in cfg.dataset:
test_list = 'annotations/instances_val2014.json'
elif '2017' in cfg.dataset:
test_list = 'annotations/instances_val2017.json'
if cfg.debug:
if not os.path.exists('output'):
os.mkdir('output')
model = YOLOv3(is_train=False)
model.build_model()
outputs = model.get_pred()
place = fluid.CUDAPlace(0) if cfg.use_gpu else fluid.CPUPlace()
exe = fluid.Executor(place)
# yapf: disable
if cfg.weights:
def if_exist(var):
return os.path.exists(os.path.join(cfg.weights, var.name))
fluid.io.load_vars(exe, cfg.weights, predicate=if_exist)
# yapf: enable
# you can save inference model by following code
# fluid.io.save_inference_model("./output/yolov3",
# feeded_var_names=['image', 'im_shape'],
# target_vars=outputs,
# executor=exe)
input_size = cfg.input_size
test_reader = reader.test(input_size, 1)
label_names, label_ids = reader.get_label_infos()
if cfg.debug:
print("Load in labels {} with ids {}".format(label_names, label_ids))
feeder = fluid.DataFeeder(place=place, feed_list=model.feeds())
def get_pred_result(boxes, scores, labels, im_id):
result = []
for box, score, label in zip(boxes, scores, labels):
x1, y1, x2, y2 = box
w = x2 - x1 + 1
h = y2 - y1 + 1
bbox = [x1, y1, w, h]
res = {
'image_id': im_id,
'category_id': label_ids[int(label)],
'bbox': list(map(float, bbox)),
'score': float(score)
}
result.append(res)
return result
dts_res = []
fetch_list = [outputs]
total_time = 0
for batch_id, batch_data in enumerate(test_reader()):
start_time = time.time()
batch_outputs = exe.run(fetch_list=[v.name for v in fetch_list],
feed=feeder.feed(batch_data),
return_numpy=False,
use_program_cache=True)
lod = batch_outputs[0].lod()[0]
nmsed_boxes = np.array(batch_outputs[0])
if nmsed_boxes.shape[1] != 6:
continue
for i in range(len(lod) - 1):
im_id = batch_data[i][1]
start = lod[i]
end = lod[i + 1]
if start == end:
continue
nmsed_box = nmsed_boxes[start:end, :]
labels = nmsed_box[:, 0]
scores = nmsed_box[:, 1]
boxes = nmsed_box[:, 2:6]
dts_res += get_pred_result(boxes, scores, labels, im_id)
end_time = time.time()
print("batch id: {}, time: {}".format(batch_id, end_time - start_time))
total_time += end_time - start_time
with open("yolov3_result.json", 'w') as outfile:
json.dump(dts_res, outfile)
print("start evaluate detection result with coco api")
coco = COCO(os.path.join(cfg.data_dir, test_list))
cocoDt = coco.loadRes("yolov3_result.json")
cocoEval = COCOeval(coco, cocoDt, 'bbox')
cocoEval.evaluate()
cocoEval.accumulate()
cocoEval.summarize()
print("evaluate done.")
print("Time per batch: {}".format(total_time / batch_id))
def train():
if cfg.debug:
fluid.default_startup_program().random_seed = 1000
fluid.default_main_program().random_seed = 1000
random.seed(0)
np.random.seed(0)
if not os.path.exists(cfg.model_save_dir):
os.makedirs(cfg.model_save_dir)
model = YOLOv3()
model.build_model()
input_size = cfg.input_size
loss = model.loss()
loss.persistable = True
devices = os.getenv("CUDA_VISIBLE_DEVICES") or ""
devices_num = len(devices.split(","))
print("Found {} CUDA devices.".format(devices_num))
learning_rate = cfg.learning_rate
boundaries = cfg.lr_steps
gamma = cfg.lr_gamma
step_num = len(cfg.lr_steps)
values = [learning_rate * (gamma**i) for i in range(step_num + 1)]
optimizer = fluid.optimizer.Momentum(
learning_rate=exponential_with_warmup_decay(
learning_rate=learning_rate,
boundaries=boundaries,
values=values,
warmup_iter=cfg.warm_up_iter,
warmup_factor=cfg.warm_up_factor),
regularization=fluid.regularizer.L2Decay(cfg.weight_decay),
momentum=cfg.momentum)
optimizer.minimize(loss)
if cfg.use_gpu:
place = fluid.CUDAPlace(0)
else:
place = fluid.CPUPlace()
exe = fluid.Executor(place)
exe.run(fluid.default_startup_program())
if cfg.pretrain:
if not os.path.exists(cfg.pretrain):
print("Pretrain weights not found: {}".format(cfg.pretrain))
def if_exist(var):
return os.path.exists(os.path.join(cfg.pretrain, var.name))
fluid.io.load_vars(exe, cfg.pretrain, predicate=if_exist)
build_strategy = fluid.BuildStrategy()
build_strategy.memory_optimize = True
build_strategy.sync_batch_norm = cfg.syncbn
compile_program = fluid.compiler.CompiledProgram(
fluid.default_main_program()).with_data_parallel(
loss_name=loss.name, build_strategy=build_strategy)
random_sizes = [cfg.input_size]
if cfg.random_shape:
random_sizes = [32 * i for i in range(10, 20)]
total_iter = cfg.max_iter - cfg.start_iter
mixup_iter = total_iter - cfg.no_mixup_iter
train_reader = reader.train(input_size,
batch_size=cfg.batch_size,
shuffle=True,
total_iter=total_iter * devices_num,
mixup_iter=mixup_iter * devices_num,
random_sizes=random_sizes,
use_multiprocessing=cfg.use_multiprocess)
py_reader = model.py_reader
py_reader.decorate_paddle_reader(train_reader)
def save_model(postfix):
model_path = os.path.join(cfg.model_save_dir, postfix)
if os.path.isdir(model_path):
shutil.rmtree(model_path)
fluid.io.save_persistables(exe, model_path)
fetch_list = [loss]
py_reader.start()
smoothed_loss = SmoothedValue()
try:
start_time = time.time()
prev_start_time = start_time
snapshot_loss = 0
snapshot_time = 0
for iter_id in range(cfg.start_iter, cfg.max_iter):
prev_start_time = start_time
start_time = time.time()
losses = exe.run(compile_program,
fetch_list=[v.name for v in fetch_list])
smoothed_loss.add_value(np.mean(np.array(losses[0])))
snapshot_loss += np.mean(np.array(losses[0]))
snapshot_time += start_time - prev_start_time
lr = np.array(
fluid.global_scope().find_var('learning_rate').get_tensor())
print("Iter {:d}, lr {:.6f}, loss {:.6f}, time {:.5f}".format(
iter_id, lr[0], smoothed_loss.get_mean_value(),
start_time - prev_start_time))
sys.stdout.flush()
if (iter_id + 1) % cfg.snapshot_iter == 0:
save_model("model_iter{}".format(iter_id))
print("Snapshot {} saved, average loss: {}, \
average time: {}".format(
iter_id + 1, snapshot_loss / float(cfg.snapshot_iter),
snapshot_time / float(cfg.snapshot_iter)))
snapshot_loss = 0
snapshot_time = 0
except fluid.core.EOFException:
py_reader.reset()
save_model('model_final')
def main():
"""
YOLOv3 trainer. See README for details.
"""
args = parse_args()
print("Setting Arguments.. : ", args)
cuda = torch.cuda.is_available() and args.use_cuda
os.makedirs(args.checkpoint_dir, exist_ok=True)
# Parse config settings
with open(args.cfg, 'r') as f:
cfg = yaml.load(f)
print("successfully loaded config file: ", cfg)
evaluate_dataset = cfg.get('EVALUATE', False)
momentum = cfg['TRAIN']['MOMENTUM']
decay = cfg['TRAIN']['DECAY']
burn_in = cfg['TRAIN']['BURN_IN']
iter_size = cfg['TRAIN']['MAXITER']
steps = eval(cfg['TRAIN']['STEPS'])
batch_size = cfg['TRAIN']['BATCHSIZE']
subdivision = cfg['TRAIN']['SUBDIVISION']
ignore_thre = cfg['TRAIN']['IGNORETHRE']
random_resize = cfg['AUGMENTATION']['RANDRESIZE']
base_lr = cfg['TRAIN']['LR'] / batch_size / subdivision
print('effective_batch_size = batch_size * iter_size = %d * %d' %
(batch_size, subdivision))
# Learning rate setup
def burnin_schedule(i):
if i < burn_in:
factor = pow(i / burn_in, 4)
elif i < steps[0]:
factor = 1.0
elif i < steps[1]:
factor = 0.1
else:
factor = 0.01
return factor
# Initiate model
model = YOLOv3(cfg['MODEL'], ignore_thre=ignore_thre)
if args.weights_path:
print("loading darknet weights....", args.weights_path)
parse_yolo_weights(model, args.weights_path)
elif args.checkpoint:
print("loading pytorch ckpt...", args.checkpoint)
state = torch.load(args.checkpoint)
if 'model_state_dict' in state.keys():
model.load_state_dict(state['model_state_dict'])
else:
model.load_state_dict(state)
if cuda:
print("using cuda")
model = model.cuda()
if args.tfboard:
print("using tfboard")
tblogger = SummaryWriter(get_tensorboard_log_path(args.dataset))
model.train()
imgsize = cfg['TRAIN']['IMGSIZE']
dataset = get_dataset_class(args.dataset)(model_type=cfg['MODEL']['TYPE'],
img_size=imgsize,
augmentation=cfg['AUGMENTATION'],
debug=args.debug)
dataloader = torch.utils.data.DataLoader(
dataset, batch_size=batch_size, shuffle=True, num_workers=args.n_cpu)
dataiterator = iter(dataloader)
evaluator = get_evaluator_class(args.dataset)(model_type=cfg['MODEL']['TYPE'],
img_size=cfg['TEST']['IMGSIZE'],
confthre=cfg['TEST']['CONFTHRE'],
nmsthre=cfg['TEST']['NMSTHRE'])
dtype = torch.cuda.FloatTensor if cuda else torch.FloatTensor
# optimizer setup
# set weight decay only on conv.weight
params_dict = dict(model.named_parameters())
params = []
for key, value in params_dict.items():
if 'conv.weight' in key:
params += [{'params': value, 'weight_decay': decay * batch_size * subdivision}]
else:
params += [{'params': value, 'weight_decay': 0.0}]
optimizer = optim.SGD(params, lr=base_lr, momentum=momentum,
dampening=0, weight_decay=decay * batch_size * subdivision)
iter_state = 0
if args.checkpoint:
if 'optimizer_state_dict' in state.keys():
optimizer.load_state_dict(state['optimizer_state_dict'])
if not evaluate_dataset:
iter_state = state['iter'] + 1
scheduler = optim.lr_scheduler.LambdaLR(optimizer, burnin_schedule)
# start training loop
start = time.time()
print('Initial iter state {}'.format(iter_state))
for iter_i in range(iter_state, iter_size + 1):
# SIGNET evaluation
if iter_i % args.eval_interval == 0 and iter_i > 0:
ap30, ar30 = evaluator.evaluate(model)
model.train()
if args.tfboard:
tblogger.add_scalar('val/SIGNETAP30', ap30, iter_i)
tblogger.add_scalar('val/SIGNETAR30', ar30, iter_i)
# subdivision loop
optimizer.zero_grad()
for _ in range(subdivision):
try:
imgs, targets, _, _ = next(dataiterator) # load a batch
except StopIteration:
dataiterator = iter(dataloader)
imgs, targets, _, _ = next(dataiterator) # load a batch
imgs = Variable(imgs.type(dtype))
targets = Variable(targets.type(dtype), requires_grad=False)
loss = model(imgs, targets)
loss.backward()
optimizer.step()
scheduler.step()
if iter_i % 10 == 0:
# logging
current_lr = scheduler.get_lr()[0] * batch_size * subdivision
print('[Iter %d/%d] [lr %f] '
'[Losses: xy %f, wh %f, conf %f, cls %f, total %f, imgsize %d]'
% (iter_i, iter_size, current_lr,
model.loss_dict['xy'], model.loss_dict['wh'],
model.loss_dict['conf'], model.loss_dict['cls'],
model.loss_dict['l2'], imgsize),
flush=True)
if args.tfboard:
tblogger.add_scalar('train/total_loss', model.loss_dict['l2'], iter_i)
# random resizing
if random_resize:
imgsize = (random.randint(0, 9) % 10 + 10) * 32
dataset.img_shape = (imgsize, imgsize)
dataset.img_size = imgsize
dataloader = torch.utils.data.DataLoader(
dataset, batch_size=batch_size, shuffle=True, num_workers=args.n_cpu)
dataiterator = iter(dataloader)
# save checkpoint
if iter_i > 0 and (iter_i % args.checkpoint_interval == 0):
torch.save({'iter': iter_i,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
},
os.path.join(args.checkpoint_dir, "snapshot"+str(iter_i)+".ckpt"))
end = time.time()
print("Training time in seconds: {}".format(end-start))
if args.tfboard:
tblogger.close()