130 lines
4.1 KiB
Python
130 lines
4.1 KiB
Python
# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import argparse
|
|
|
|
PARSER = argparse.ArgumentParser(description="VNet")
|
|
|
|
PARSER.add_argument('--exec_mode',
|
|
choices=['train', 'predict', 'train_and_predict', 'train_and_evaluate'],
|
|
required=True,
|
|
type=str)
|
|
|
|
PARSER.add_argument('--data_normalization',
|
|
choices=['zscore'],
|
|
default='zscore',
|
|
type=str)
|
|
|
|
PARSER.add_argument('--activation',
|
|
choices=['relu'],
|
|
default='relu',
|
|
type=str)
|
|
|
|
PARSER.add_argument('--resize_interpolator',
|
|
choices=['linear'],
|
|
default='linear',
|
|
type=str)
|
|
|
|
PARSER.add_argument('--loss',
|
|
choices=['dice'],
|
|
default='dice',
|
|
type=str)
|
|
|
|
PARSER.add_argument('--normalization_layer',
|
|
choices=['batchnorm'],
|
|
default='batchnorm',
|
|
type=str)
|
|
|
|
PARSER.add_argument('--pooling',
|
|
choices=['conv_pool'],
|
|
default='conv_pool',
|
|
type=str)
|
|
|
|
PARSER.add_argument('--upsampling',
|
|
choices=['transposed_conv'],
|
|
default='transposed_conv',
|
|
type=str)
|
|
|
|
PARSER.add_argument('--seed',
|
|
default=0,
|
|
type=int)
|
|
|
|
PARSER.add_argument('--input_shape', nargs='+', type=int, default=[32, 32, 32])
|
|
PARSER.add_argument('--upscale_blocks', nargs='+', type=int, default=[3, 3])
|
|
PARSER.add_argument('--downscale_blocks', nargs='+', type=int, default=[3, 3, 3])
|
|
|
|
PARSER.add_argument('--convolution_size',
|
|
choices=[3, 5],
|
|
default=3,
|
|
type=int)
|
|
|
|
PARSER.add_argument('--batch_size',
|
|
required=True,
|
|
type=int)
|
|
|
|
PARSER.add_argument('--log_every',
|
|
default=10,
|
|
type=int)
|
|
|
|
PARSER.add_argument('--warmup_steps',
|
|
default=200,
|
|
type=int)
|
|
|
|
PARSER.add_argument('--train_epochs',
|
|
default=1,
|
|
type=int)
|
|
|
|
PARSER.add_argument('--optimizer',
|
|
choices=['rmsprop'],
|
|
default='rmsprop',
|
|
type=str)
|
|
|
|
PARSER.add_argument('--gradient_clipping',
|
|
choices=['global_norm'],
|
|
default='global_norm',
|
|
type=str)
|
|
|
|
PARSER.add_argument('--base_lr',
|
|
default=0.0001,
|
|
type=float)
|
|
|
|
PARSER.add_argument('--momentum',
|
|
default=0.0,
|
|
type=float)
|
|
|
|
PARSER.add_argument('--train_split',
|
|
default=1.0,
|
|
type=float)
|
|
|
|
PARSER.add_argument('--split_seed',
|
|
default=0,
|
|
type=int)
|
|
|
|
PARSER.add_argument('--model_dir',
|
|
required=True,
|
|
type=str)
|
|
|
|
PARSER.add_argument('--log_dir',
|
|
default=None,
|
|
type=str)
|
|
|
|
PARSER.add_argument('--data_dir',
|
|
required=True,
|
|
type=str)
|
|
|
|
PARSER.add_argument('--benchmark', dest='benchmark', action='store_true', default=False)
|
|
PARSER.add_argument('--use_amp', '--amp', dest='use_amp', action='store_true', default=False)
|
|
PARSER.add_argument('--use_xla', '--xla', dest='use_xla', action='store_true', default=False)
|
|
PARSER.add_argument('--augment', dest='augment', action='store_true', default=False)
|