return given network
(args, net, use_gpu=True, gpu_device = 0, distribution = True)
| 20 | device = torch.device('cuda', args.gpu_device) |
| 21 | |
| 22 | def get_network(args, net, use_gpu=True, gpu_device = 0, distribution = True): |
| 23 | """ return given network |
| 24 | """ |
| 25 | |
| 26 | if net == 'sam2': |
| 27 | from sam2_train.build_sam import build_sam2_video_predictor |
| 28 | |
| 29 | sam2_checkpoint = args.sam_ckpt |
| 30 | model_cfg = args.sam_config |
| 31 | |
| 32 | net = build_sam2_video_predictor(config_file=model_cfg, ckpt_path=sam2_checkpoint, mode=None) |
| 33 | else: |
| 34 | print('the network name you have entered is not supported yet') |
| 35 | sys.exit() |
| 36 | |
| 37 | if use_gpu: |
| 38 | net = net.to(device=gpu_device) |
| 39 | |
| 40 | return net |
| 41 | |
| 42 | def create_logger(log_dir, phase='train'): |
| 43 | time_str = time.strftime('%Y-%m-%d-%H-%M') |
nothing calls this directly
no test coverage detected