(args, net: nn.Module, optimizer1, optimizer2, train_loader,
epoch)
| 48 | |
| 49 | |
| 50 | def train_sam(args, net: nn.Module, optimizer1, optimizer2, train_loader, |
| 51 | epoch): |
| 52 | hard = 0 |
| 53 | epoch_loss = 0 |
| 54 | epoch_prompt_loss = 0 |
| 55 | epoch_non_prompt_loss = 0 |
| 56 | ind = 0 |
| 57 | # train mode |
| 58 | net.train() |
| 59 | if optimizer1 is not None: |
| 60 | optimizer1.zero_grad() |
| 61 | if optimizer2 is not None: |
| 62 | optimizer2.zero_grad() |
| 63 | video_length = args.video_length |
| 64 | |
| 65 | GPUdevice = torch.device('cuda:' + str(args.gpu_device)) |
| 66 | prompt = args.prompt |
| 67 | prompt_freq = args.prompt_freq |
| 68 | |
| 69 | lossfunc = criterion_G |
| 70 | # lossfunc = paper_loss#.to(dtype=torch.bfloat16, device=GPUdevice) |
| 71 | |
| 72 | with tqdm(total=len(train_loader), desc=f'Epoch {epoch}', unit='img') as pbar: |
| 73 | for pack in train_loader: |
| 74 | torch.cuda.empty_cache() |
| 75 | imgs_tensor = pack['image'] |
| 76 | mask_dict = pack['label'] |
| 77 | if prompt == 'click': |
| 78 | pt_dict = pack['pt'] |
| 79 | point_labels_dict = pack['p_label'] |
| 80 | elif prompt == 'bbox': |
| 81 | bbox_dict = pack['bbox'] |
| 82 | imgs_tensor = imgs_tensor.squeeze(0) |
| 83 | imgs_tensor = imgs_tensor.to(dtype = torch.float32, device = GPUdevice) |
| 84 | |
| 85 | train_state = net.train_init_state(imgs_tensor=imgs_tensor) |
| 86 | prompt_frame_id = list(range(0, video_length, prompt_freq)) |
| 87 | obj_list = [] |
| 88 | for id in prompt_frame_id: |
| 89 | obj_list += list(mask_dict[id].keys()) |
| 90 | obj_list = list(set(obj_list)) |
| 91 | if len(obj_list) == 0: |
| 92 | continue |
| 93 | |
| 94 | name = pack['image_meta_dict']['filename_or_obj'] |
| 95 | # reverse = np.random.rand() > 0.5 |
| 96 | |
| 97 | with torch.cuda.amp.autocast(): |
| 98 | for id in prompt_frame_id: |
| 99 | for ann_obj_id in obj_list: |
| 100 | try: |
| 101 | if prompt == 'click': |
| 102 | points = pt_dict[id][ann_obj_id].to(device=GPUdevice) |
| 103 | labels = point_labels_dict[id][ann_obj_id].to(device=GPUdevice) |
| 104 | _, _, _ = net.train_add_new_points( |
| 105 | inference_state=train_state, |
| 106 | frame_idx=id, |
| 107 | obj_id=ann_obj_id, |
nothing calls this directly
no test coverage detected