return given network
(args, net, use_gpu=True, gpu_device = 0, distribution = True)
| 41 | |
| 42 | |
| 43 | def get_network(args, net, use_gpu=True, gpu_device = 0, distribution = True): |
| 44 | """ return given network |
| 45 | """ |
| 46 | |
| 47 | |
| 48 | if net == 'sam2': |
| 49 | from sam2_train.build_sam import build_sam2 |
| 50 | from sam2_train.sam2_image_predictor import SAM2ImagePredictor |
| 51 | torch.autocast(device_type="cuda", dtype=torch.bfloat16).__enter__() |
| 52 | if torch.cuda.get_device_properties(0).major >= 8: |
| 53 | # turn on tfloat32 for Ampere GPUs (https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices) |
| 54 | torch.backends.cuda.matmul.allow_tf32 = True |
| 55 | torch.backends.cudnn.allow_tf32 = True |
| 56 | net = build_sam2(args.sam_config, args.sam_ckpt, device="cuda") |
| 57 | |
| 58 | |
| 59 | else: |
| 60 | print('the network name you have entered is not supported yet') |
| 61 | sys.exit() |
| 62 | |
| 63 | if use_gpu: |
| 64 | #net = net.cuda(device = gpu_device) |
| 65 | if distribution != 'none': |
| 66 | net = torch.nn.DataParallel(net,device_ids=[int(id) for id in args.distributed.split(',')]) |
| 67 | net = net.to(device=gpu_device) |
| 68 | else: |
| 69 | net = net.to(device=gpu_device) |
| 70 | |
| 71 | return net |
| 72 | |
| 73 | |
| 74 | @torch.no_grad() |
no test coverage detected