is_horovod = cfg.TRAINER == 'horovod'
if is_horovod:
hvd.init()
logger.info("Horovod Rank={}, Size={}".format(
hvd.rank(), hvd.size()))
if not is_horovod or hvd.rank() == 0:
logger.set_logger_dir(args.logdir, 'd')
log_launch_config(args.log_full_git_diff)
finalize_configs(is_training=True)
if cfg.TRAIN.SEED:
tf.set_random_seed(cfg.TRAIN.SEED)
fix_rng_seed(cfg.TRAIN.SEED * hvd.rank())
np.random.seed(cfg.TRAIN.SEED)
images_per_step = cfg.TRAIN.NUM_GPUS * cfg.TRAIN.BATCH_SIZE_PER_GPU
steps_per_epoch = args.images_per_epoch // images_per_step
batch_size_lr_factor = images_per_step # The LR is defined for bs=1 and then scaled linearly with the batch size
base_lr_adjusted_for_bs = cfg.TRAIN.BASE_LR * batch_size_lr_factor
# Warmup LR schedule is step based
warmup_start_step = 0
warmup_end_step = cfg.TRAIN.WARMUP_STEPS
warmup_start_lr = cfg.TRAIN.WARMUP_INIT_LR * 8
warmup_end_lr = base_lr_adjusted_for_bs
warmup_schedule = [(warmup_start_step, warmup_start_lr),
(warmup_end_step, warmup_end_lr)]
parser.add_argument('--save', default='train_log/', help='Saving directory')
parser.add_argument('--mode', default='none', help='Additional mode of resnet')
parser.add_argument('--types', type=int, default=16)
parser.add_argument('--threshold', type=float, default=0.5)
parser.add_argument('--pathology', default='All')
parser.add_argument('--batch', type=int, default=64)
parser.add_argument('--shape', type=int, default=256)
args = parser.parse_args()
if args.seed:
os.environ['PYTHONHASHSEED']=str(args.seed)
random.seed(args.seed)
np.random.seed(args.seed)
fix_rng_seed(args.seed)
tf.random.set_random_seed(args.seed)
if args.gpus:
# os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpus
model = Model(args=args)
if args.eval:
ds_valid = Vinmec(folder=args.data,
is_train='valid',
fname='valid.csv',
types=args.types,
pathology=args.pathology,
resize=int(args.shape))