Skip to content

Commit

Permalink
bin/train.py: Update for updated named_subparser API.
Browse files Browse the repository at this point in the history
  • Loading branch information
de-vri-es committed Feb 25, 2018
1 parent 7506b0e commit 6bfedce
Showing 1 changed file with 6 additions and 8 deletions.
14 changes: 6 additions & 8 deletions keras_retinanet/bin/train.py
Original file line number Diff line number Diff line change
Expand Up @@ -175,8 +175,8 @@ def check_args(parsed_args, dataset_plugins):

# let the selected dataset check args too
dataset_plugins[parsed_args.dataset.name].check_args(parsed_args.dataset.args)
if parsed_args.evaluate:
dataset_plugins[parsed_args.evaluate.name].check_args(parsed_args.evaluate.args)
if parsed_args.validate:
dataset_plugins[parsed_args.validate.name].check_args(parsed_args.validate.args)

return parsed_args

Expand Down Expand Up @@ -210,10 +210,8 @@ def parse_args(args, dataset_plugins):
parser.add_argument('--tensorboard-dir', help='Log directory for Tensorboard output', default='./logs')
parser.add_argument('--no-snapshots', help='Disable saving snapshots.', dest='snapshots', action='store_false')

train_dataset = NamedSubparser('--dataset', required=True)
validation_dataset = NamedSubparser('--evaluate', required=False)
parser.add_named_subparser(train_dataset)
parser.add_named_subparser(validation_dataset)
train_dataset = parser.add_named_subparser(['--dataset'], required=True)
validation_dataset = parser.add_named_subparser(['--validate'])

# let all plugins register their arguments.
for name, plugin in dataset_plugins.items():
Expand Down Expand Up @@ -259,8 +257,8 @@ def main(args=None):

# make the validation data generator
validation_generator = None
if args.evaluate:
validation_generator = create_generator(arg.evaluate, dataset_plugins, batch_size=args.batch_size)
if args.validate:
validation_generator = create_generator(arg.validate, dataset_plugins, batch_size=args.batch_size)

if 'resnet' in args.backbone:
from ..models.resnet import resnet_retinanet as retinanet, custom_objects, download_imagenet
Expand Down

0 comments on commit 6bfedce

Please sign in to comment.