def main():
parser = argparse.ArgumentParser()
parser.add_argument('--model',
choices=('faster_rcnn_fpn_resnet50',
'faster_rcnn_fpn_resnet101'),
default='faster_rcnn_fpn_resnet50')
parser.add_argument('--gpu', type=int, default=-1)
parser.add_argument('--pretrained-model', default='coco')
parser.add_argument('image')
args = parser.parse_args()
if args.model == 'faster_rcnn_fpn_resnet50':
model = FasterRCNNFPNResNet50(n_fg_class=len(coco_bbox_label_names),
pretrained_model=args.pretrained_model)
elif args.model == 'faster_rcnn_fpn_resnet101':
model = FasterRCNNFPNResNet101(n_fg_class=len(coco_bbox_label_names),
pretrained_model=args.pretrained_model)
if args.gpu >= 0:
chainer.cuda.get_device_from_id(args.gpu).use()
model.to_gpu()
img = utils.read_image(args.image)
bboxes, labels, scores = model.predict([img])
bbox = bboxes[0]
label = labels[0]
score = scores[0]
vis_bbox(img, bbox, label, score, label_names=coco_bbox_label_names)
plt.show()
def __init__(self, gpu, model, nms_thresh= 0.45, score_thresh=0.6):
self.gpu = gpu
if model == 'yolo_v2_tiny':
self.model = YOLOv2Tiny(
n_fg_class=len(voc_bbox_label_names),
pretrained_model='voc0712')
elif model == 'yolo_v3':
self.model = YOLOv3(
n_fg_class=len(voc_bbox_label_names),
pretrained_model='voc0712')
elif model == 'ssd300':
self.model = SSD300(
n_fg_class=len(voc_bbox_label_names),
pretrained_model='voc0712')
elif model == 'ssd512':
self.model = SSD512(
n_fg_class=len(voc_bbox_label_names),
pretrained_model='voc0712')
elif model == 'fasterrcnnvgg16':
self.model = FasterRCNNVGG16(
n_fg_class=len(coco_bbox_label_names),
pretrained_model='voc0712')
elif model == 'fasterrcnnfpnresnet50':
self.model = FasterRCNNFPNResNet50(
n_fg_class=len(coco_bbox_label_names),
pretrained_model='coco')
elif model == 'fasterrcnnfpnresnet101':
self.model = FasterRCNNFPNResNet101(
n_fg_class=len(coco_bbox_label_names),
pretrained_model='coco')
else:
self.model = YOLOv2(
n_fg_class=len(voc_bbox_label_names),
pretrained_model='voc0712')
#self.model.nms_thresh = nms_thresh
#self.model.score_thresh = score_thresh
if self.gpu >= 0:
chainer.cuda.get_device_from_id(self.gpu).use()
self.model.to_gpu()
def main():
parser = argparse.ArgumentParser()
parser.add_argument(
'--model',
choices=('faster_rcnn_fpn_resnet50', 'faster_rcnn_fpn_resnet101'),
default='faster_rcnn_fpn_resnet50')
parser.add_argument('--batchsize', type=int, default=16)
parser.add_argument('--iteration', type=int, default=90000)
parser.add_argument('--step', type=int, nargs='*', default=[60000, 80000])
parser.add_argument('--out', default='result')
parser.add_argument('--resume')
args = parser.parse_args()
# https://docs.chainer.org/en/stable/chainermn/tutorial/tips_faqs.html#using-multiprocessiterator
if hasattr(multiprocessing, 'set_start_method'):
multiprocessing.set_start_method('forkserver')
p = multiprocessing.Process()
p.start()
p.join()
comm = chainermn.create_communicator()
device = comm.intra_rank
if args.model == 'faster_rcnn_fpn_resnet50':
model = FasterRCNNFPNResNet50(
n_fg_class=len(coco_bbox_label_names), pretrained_model='imagenet')
elif args.model == 'faster_rcnn_fpn_resnet101':
model = FasterRCNNFPNResNet101(
n_fg_class=len(coco_bbox_label_names), pretrained_model='imagenet')
model.use_preset('evaluate')
train_chain = TrainChain(model)
chainer.cuda.get_device_from_id(device).use()
train_chain.to_gpu()
train = TransformDataset(
COCOBboxDataset(year='2017', split='train'),
('img', 'bbox', 'label'), transform)
if comm.rank == 0:
indices = np.arange(len(train))
else:
indices = None
indices = chainermn.scatter_dataset(indices, comm, shuffle=True)
train = train.slice[indices]
train_iter = chainer.iterators.MultithreadIterator(
train, args.batchsize // comm.size)
optimizer = chainermn.create_multi_node_optimizer(
chainer.optimizers.MomentumSGD(), comm)
optimizer.setup(train_chain)
optimizer.add_hook(WeightDecay(0.0001))
model.extractor.base.conv1.disable_update()
model.extractor.base.res2.disable_update()
for link in model.links():
if isinstance(link, L.BatchNormalization):
link.disable_update()
updater = training.updaters.StandardUpdater(
train_iter, optimizer, converter=converter, device=device)
trainer = training.Trainer(
updater, (args.iteration * 16 / args.batchsize, 'iteration'), args.out)
@make_shift('lr')
def lr_schedule(trainer):
base_lr = 0.02 * args.batchsize / 16
warm_up_duration = 500
warm_up_rate = 1 / 3
iteration = trainer.updater.iteration
if iteration < warm_up_duration:
rate = warm_up_rate \
+ (1 - warm_up_rate) * iteration / warm_up_duration
else:
rate = 1
for step in args.step:
if iteration >= step * 16 / args.batchsize:
rate *= 0.1
return base_lr * rate
trainer.extend(lr_schedule)
if comm.rank == 0:
log_interval = 10, 'iteration'
trainer.extend(extensions.LogReport(trigger=log_interval))
trainer.extend(extensions.observe_lr(), trigger=log_interval)
trainer.extend(extensions.PrintReport(
['epoch', 'iteration', 'lr', 'main/loss',
'main/loss/rpn/loc', 'main/loss/rpn/conf',
'main/loss/head/loc', 'main/loss/head/conf']),
trigger=log_interval)
trainer.extend(extensions.ProgressBar(update_interval=10))
trainer.extend(extensions.snapshot(), trigger=(10000, 'iteration'))
trainer.extend(
extensions.snapshot_object(
model, 'model_iter_{.updater.iteration}'),
trigger=(90000 * 16 / args.batchsize, 'iteration'))
if args.resume:
serializers.load_npz(args.resume, trainer, strict=False)
trainer.run()