()
| 474 | |
| 475 | |
| 476 | def main(): |
| 477 | args = parse_args() |
| 478 | logger.info('Called with argument: %s' % args) |
| 479 | #ctx = [mx.gpu(int(i)) for i in args.gpus.split(',')] |
| 480 | ctx = [] |
| 481 | cvd = os.environ['CUDA_VISIBLE_DEVICES'].strip() |
| 482 | if len(cvd) > 0: |
| 483 | for i in range(len(cvd.split(','))): |
| 484 | ctx.append(mx.gpu(i)) |
| 485 | if len(ctx) == 0: |
| 486 | ctx = [mx.cpu()] |
| 487 | print('use cpu') |
| 488 | else: |
| 489 | print('gpu num:', len(ctx)) |
| 490 | train_net(args, |
| 491 | ctx, |
| 492 | args.pretrained, |
| 493 | args.pretrained_epoch, |
| 494 | args.prefix, |
| 495 | args.begin_epoch, |
| 496 | args.end_epoch, |
| 497 | lr=args.lr, |
| 498 | lr_step=args.lr_step) |
| 499 | |
| 500 | |
| 501 | if __name__ == '__main__': |
no test coverage detected