| 362 | |
| 363 | |
| 364 | def main(): |
| 365 | parser = argparse.ArgumentParser(description='U-Net model training') |
| 366 | parser.add_argument('-dataset', |
| 367 | default='em', |
| 368 | help='Dataset.') |
| 369 | parser.add_argument('-learning_rate', |
| 370 | default=0.0001, |
| 371 | help='Learning Rate') |
| 372 | parser.add_argument('-penal_cst', |
| 373 | default=0.0, |
| 374 | help='regularization constant') |
| 375 | parser.add_argument('--num_epochs', |
| 376 | '-ne', |
| 377 | type=int, |
| 378 | default=750, |
| 379 | help='Optional. Int to indicate the max' |
| 380 | 'number of epochs.') |
| 381 | parser.add_argument('-max_patience', |
| 382 | type=int, |
| 383 | default=100, |
| 384 | help='Max patience') |
| 385 | parser.add_argument('-batch_size', |
| 386 | type=int, |
| 387 | nargs='+', |
| 388 | default=[5, 5, 1], |
| 389 | help='Batch size [train, val, test]. Default: -batch_size 5 5 1') |
| 390 | parser.add_argument('-data_augmentation', |
| 391 | type=json.loads, |
| 392 | default={'rotation_range':25, |
| 393 | 'shear_range':0.41, |
| 394 | 'horizontal_flip':True, |
| 395 | 'vertical_flip':True, |
| 396 | 'fill_mode':'reflect', |
| 397 | 'spline_warp':True, |
| 398 | 'warp_sigma':10, |
| 399 | 'warp_grid_size':3, |
| 400 | 'crop_size': (224, 224)}, |
| 401 | help='use data augmentation') |
| 402 | parser.add_argument('-early_stop_class', |
| 403 | type=int, |
| 404 | default=None, |
| 405 | help='class to early stop on') |
| 406 | parser.add_argument('-train_from_0_255', |
| 407 | type=bool, |
| 408 | default=False, |
| 409 | help='Whether to train from images within 0-255 range') |
| 410 | args = parser.parse_args() |
| 411 | |
| 412 | train(args.dataset, float(args.learning_rate), |
| 413 | float(args.penal_cst), int(args.num_epochs), int(args.max_patience), |
| 414 | data_augmentation=args.data_augmentation, batch_size=args.batch_size, |
| 415 | early_stop_class=args.early_stop_class, savepath=SAVEPATH, |
| 416 | train_from_0_255=args.train_from_0_255, loadpath=LOADPATH) |
| 417 | |
| 418 | if __name__ == "__main__": |
| 419 | main() |