import os import time from tqdm import tqdm import cv2 import numpy as np import torch import torch.cuda.amp as amp import torch.distributed as dist import torch.nn.functional as F import wandb from loguru import logger from utils.wandb_utils import fix_metric_names # keep from utils.dataset import tokenize # keep from utils.misc import ( AverageMeter, ProgressMeter, concat_all_gather, trainMetricGPU, # NOTE: no longer used (ins mask metrics) get_seg_image # NOTE: no longer used (ins mask vis) ) from utils.grasp_eval import ( detect_grasps, calculate_iou, calculate_max_iou, calculate_jacquard_index, visualization, ) def _all_ranks_finite_loss(loss, device): local_finite = torch.tensor( [1 if torch.isfinite(loss.detach()).item() else 0], device=device, dtype=torch.int32, ) if dist.is_available() and dist.is_initialized(): dist.all_reduce(local_finite, op=dist.ReduceOp.MIN) return bool(local_finite.item()) def _format_tensor_stats(name, tensor): if not torch.is_tensor(tensor): return f"{name}=non_tensor" with torch.no_grad(): data = tensor.detach() finite = torch.isfinite(data) finite_ratio = finite.float().mean().item() if finite.any(): finite_data = data[finite] min_value = finite_data.min().item() max_value = finite_data.max().item() mean_value = finite_data.float().mean().item() return ( f"{name}: finite={finite_ratio:.4f}, " f"min={min_value:.4g}, max={max_value:.4g}, mean={mean_value:.4g}" ) return f"{name}: finite={finite_ratio:.4f}" def train_with_grasp(train_loader, model, optimizer, scheduler, scaler, epoch, args): rank = dist.get_rank() if dist.is_available() and dist.is_initialized() else 0 batch_time = AverageMeter('Batch', ':2.2f') data_time = AverageMeter('Data', ':2.2f') lr = AverageMeter('Lr', ':1.6f') loss_meter = AverageMeter('Loss', ':2.4f') qua_loss_metter = AverageMeter('Loss_qua', ':2.4f') sin_loss_metter = AverageMeter('Loss_sin', ':2.4f') cos_loss_metter = AverageMeter('Loss_cos', ':2.4f') wid_loss_metter = AverageMeter('Loss_wid', ':2.4f') # WebDataset / IterableDataset may not support len(train_loader) steps_per_epoch = getattr(args, 'steps_per_epoch', None) if steps_per_epoch is None: try: steps_per_epoch = len(train_loader) except TypeError: raise ValueError( "train_loader has no length. Please set `steps_per_epoch` in yaml." ) progress = ProgressMeter( steps_per_epoch, [ batch_time, data_time, lr, loss_meter, qua_loss_metter, sin_loss_metter, cos_loss_metter, wid_loss_metter, ], prefix="Training: Epoch=[{}/{}] ".format(epoch, args.epochs) ) model.train() time.sleep(2) end = time.time() debug_steps = max(0, int(getattr(args, "ddp_debug_steps", 0))) debug_every = max(0, int(getattr(args, "ddp_debug_every", 0))) loader_iter = iter(train_loader) for i in range(steps_per_epoch): trace_step = i < debug_steps or (debug_every > 0 and ((i + 1) % debug_every == 0)) if trace_step: logger.info(f"[Rank {rank}] step {i + 1}/{steps_per_epoch}: waiting for next batch") try: data = next(loader_iter) except StopIteration: logger.warning(f"[Rank {rank}] dataloader exhausted early at step {i + 1}/{steps_per_epoch}") break if trace_step: batch_size = data["img"].size(0) if "img" in data and hasattr(data["img"], "size") else "unknown" sample_ids = data.get("sample_id", []) sample_preview = sample_ids[:2] if isinstance(sample_ids, list) else sample_ids logger.info( f"[Rank {rank}] step {i + 1}/{steps_per_epoch}: dataloader batch received " f"(batch_size={batch_size}, sample_preview={sample_preview})" ) try: image = data["img"] text = data["word_vec"] if not getattr(args, 'use_raw_text', False) else data["sentence"] # ====== no instance mask ====== grasp_qua_mask = data["grasp_masks"]["qua"] grasp_sin_mask = data["grasp_masks"]["sin"] grasp_cos_mask = data["grasp_masks"]["cos"] grasp_wid_mask = data["grasp_masks"]["wid"] data_time.update(time.time() - end) image = image.cuda(non_blocking=True) if isinstance(text, torch.Tensor): text = text.cuda(non_blocking=True) grasp_qua_mask = grasp_qua_mask.cuda(non_blocking=True).unsqueeze(1) grasp_sin_mask = grasp_sin_mask.cuda(non_blocking=True).unsqueeze(1) grasp_cos_mask = grasp_cos_mask.cuda(non_blocking=True).unsqueeze(1) grasp_wid_mask = grasp_wid_mask.cuda(non_blocking=True).unsqueeze(1) if trace_step: logger.info(f"[Rank {rank}] step {i + 1}/{steps_per_epoch}: forward start") # forward with amp.autocast(): pred, target, loss, loss_dict = model( image, text, grasp_qua_mask, grasp_sin_mask, grasp_cos_mask, grasp_wid_mask ) if not _all_ranks_finite_loss(loss, image.device): optimizer.zero_grad(set_to_none=True) sample_ids = data.get("sample_id", []) sample_preview = sample_ids[:2] if isinstance(sample_ids, list) else sample_ids logger.warning( "[Rank {}] non-finite loss at epoch {} step {}; skipping optimizer step. " "loss_dict={} sample_preview={} {} {} {} {}".format( rank, epoch, i + 1, loss_dict, sample_preview, _format_tensor_stats("pred_qua", pred[0]), _format_tensor_stats("pred_sin", pred[1]), _format_tensor_stats("pred_cos", pred[2]), _format_tensor_stats("pred_wid", pred[3]), ) ) batch_time.update(time.time() - end) end = time.time() continue if trace_step: logger.info( f"[Rank {rank}] step {i + 1}/{steps_per_epoch}: forward done " f"(loss={float(loss.detach().item()):.6f})" ) logger.info(f"[Rank {rank}] step {i + 1}/{steps_per_epoch}: backward start") # backward optimizer.zero_grad() scaler.scale(loss).backward() if trace_step: logger.info(f"[Rank {rank}] step {i + 1}/{steps_per_epoch}: backward done") if args.max_norm: scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_( model.parameters(), args.max_norm, error_if_nonfinite=False, ) scaler.step(optimizer) scaler.update() if trace_step: logger.info(f"[Rank {rank}] step {i + 1}/{steps_per_epoch}: optimizer step done") # scheduler step if hasattr(scheduler, '_scheduler_type'): if scheduler._scheduler_type == 'OneCycleLR': scheduler.step() elif hasattr(scheduler, '_warmup_steps') and scheduler._step_count < scheduler._warmup_steps: scheduler.step() elif hasattr(scheduler, '__class__') and scheduler.__class__.__name__ == 'OneCycleLR': scheduler.step() if trace_step: logger.info(f"[Rank {rank}] step {i + 1}/{steps_per_epoch}: all_reduce start") # distributed reduce loss_detached = loss.detach() dist.all_reduce(loss_detached) loss_detached = loss_detached / dist.get_world_size() if trace_step: logger.info(f"[Rank {rank}] step {i + 1}/{steps_per_epoch}: all_reduce done") except Exception: logger.exception(f"[Rank {rank}] step {i + 1}/{steps_per_epoch} failed before collective finished") raise loss_meter.update(loss_detached.item(), image.size(0)) qua_loss_metter.update(loss_dict["m_qua"], image.size(0)) sin_loss_metter.update(loss_dict["m_sin"], image.size(0)) cos_loss_metter.update(loss_dict["m_cos"], image.size(0)) wid_loss_metter.update(loss_dict["m_wid"], image.size(0)) lr.update(scheduler.get_last_lr()[-1] if scheduler is not None else optimizer.param_groups[0]['lr']) batch_time.update(time.time() - end) end = time.time() if (i + 1) % args.print_freq == 0 or (i + 1) == steps_per_epoch: progress.display(i + 1) if dist.get_rank() in [-1, 0]: metrics = { "time/batch": batch_time.val, "time/data": data_time.val, "training/lr": lr.val, "training/loss": loss_meter.val, "training/loss_qua": qua_loss_metter.val, "training/loss_sin": sin_loss_metter.val, "training/loss_cos": cos_loss_metter.val, "training/loss_wid": wid_loss_metter.val, } if hasattr(optimizer, 'param_groups'): for idx, param_group in enumerate(optimizer.param_groups): metrics[f"lr/param_group_{idx}"] = param_group['lr'] global_step = (epoch - 1) * steps_per_epoch + i wandb.log(metrics, step=global_step) @torch.no_grad() def validate_with_grasp(val_loader, model, epoch, args): def inverse(img, mat, w, h): img = np.asarray(img) if img.dtype != np.float32: img = img.astype(np.float32) if torch.is_tensor(mat): mat = mat.detach().cpu().numpy() mat = np.asarray(mat, dtype=np.float32) if torch.is_tensor(w): w = int(w.item()) else: w = int(w) if torch.is_tensor(h): h = int(h.item()) else: h = int(h) inv_img = cv2.warpAffine( img, mat, (w, h), flags=cv2.INTER_CUBIC, borderValue=0. ) return inv_img model.eval() time.sleep(2) num_grasps = [1, 5] num_correct_grasps = [0, 0] num_total_grasps = [0, 0] val_steps_per_epoch = getattr(args, 'val_steps_per_epoch', None) pbar = tqdm(val_loader, total=val_steps_per_epoch) for step, data in enumerate(pbar): # stop manually for iterable-style dataloader if val_steps_per_epoch is not None and step >= val_steps_per_epoch: break image = data["img"] text = data["word_vec"] if not getattr(args, 'use_raw_text', False) else data["sentence"] # ====== no instance mask ====== grasp_qua_mask = data["grasp_masks"]["qua"] grasp_sin_mask = data["grasp_masks"]["sin"] grasp_cos_mask = data["grasp_masks"]["cos"] grasp_wid_mask = data["grasp_masks"]["wid"] inverse_matrix = data["inverse"] ori_sizes = data["ori_size"] grasp_targets = data["grasps"] image = image.cuda(non_blocking=True) if isinstance(text, torch.Tensor): text = text.cuda(non_blocking=True) grasp_qua_mask = grasp_qua_mask.cuda(non_blocking=True).unsqueeze(1) grasp_sin_mask = grasp_sin_mask.cuda(non_blocking=True).unsqueeze(1) grasp_cos_mask = grasp_cos_mask.cuda(non_blocking=True).unsqueeze(1) grasp_wid_mask = grasp_wid_mask.cuda(non_blocking=True).unsqueeze(1) with amp.autocast(): pred, target = model( image, text, grasp_qua_mask, grasp_sin_mask, grasp_cos_mask, grasp_wid_mask ) # predictions (no ins mask) grasp_qua_mask_preds = pred[0] grasp_sin_mask_preds = pred[1] grasp_cos_mask_preds = pred[2] grasp_wid_mask_preds = pred[3] # targets (no ins mask) grasp_qua_mask_targets = target[0] grasp_sin_mask_targets = target[1] grasp_cos_mask_targets = target[2] grasp_wid_mask_targets = target[3] # sigmoid where needed (qua, wid) grasp_qua_mask_preds = torch.sigmoid(grasp_qua_mask_preds) grasp_wid_mask_preds = torch.sigmoid(grasp_wid_mask_preds) # resize preds to input image size if grasp_qua_mask_preds.shape[-2:] != image.shape[-2:]: grasp_qua_mask_preds = F.interpolate( grasp_qua_mask_preds, size=image.shape[-2:], mode='bicubic', align_corners=True ).squeeze(1) grasp_sin_mask_preds = F.interpolate( grasp_sin_mask_preds, size=image.shape[-2:], mode='bicubic', align_corners=True ).squeeze(1) grasp_cos_mask_preds = F.interpolate( grasp_cos_mask_preds, size=image.shape[-2:], mode='bicubic', align_corners=True ).squeeze(1) grasp_wid_mask_preds = F.interpolate( grasp_wid_mask_preds, size=image.shape[-2:], mode='bicubic', align_corners=True ).squeeze(1) else: grasp_qua_mask_preds = grasp_qua_mask_preds.squeeze(1) grasp_sin_mask_preds = grasp_sin_mask_preds.squeeze(1) grasp_cos_mask_preds = grasp_cos_mask_preds.squeeze(1) grasp_wid_mask_preds = grasp_wid_mask_preds.squeeze(1) # iterate over batch for idx in range(grasp_qua_mask_preds.shape[0]): inv_mat = inverse_matrix[idx] ori_size = ori_sizes[idx] h, w = ori_size grasp_qua_mask_pred = grasp_qua_mask_preds[idx].squeeze().cpu().numpy() grasp_sin_mask_pred = grasp_sin_mask_preds[idx].squeeze().cpu().numpy() grasp_cos_mask_pred = grasp_cos_mask_preds[idx].squeeze().cpu().numpy() grasp_wid_mask_pred = grasp_wid_mask_preds[idx].squeeze().cpu().numpy() grasp_target = grasp_targets[idx] grasp_qua_mask_target = grasp_qua_mask_targets[idx].squeeze().cpu().numpy() grasp_sin_mask_target = grasp_sin_mask_targets[idx].squeeze().cpu().numpy() grasp_cos_mask_target = grasp_cos_mask_targets[idx].squeeze().cpu().numpy() grasp_wid_mask_target = grasp_wid_mask_targets[idx].squeeze().cpu().numpy() # inverse to original size grasp_qua_mask_pred = inverse(grasp_qua_mask_pred, inv_mat, w, h) grasp_sin_mask_pred = inverse(grasp_sin_mask_pred, inv_mat, w, h) grasp_cos_mask_pred = inverse(grasp_cos_mask_pred, inv_mat, w, h) grasp_wid_mask_pred = inverse(grasp_wid_mask_pred, inv_mat, w, h) grasp_qua_mask_target = inverse(grasp_qua_mask_target, inv_mat, w, h) grasp_sin_mask_target = inverse(grasp_sin_mask_target, inv_mat, w, h) grasp_cos_mask_target = inverse(grasp_cos_mask_target, inv_mat, w, h) grasp_wid_mask_target = inverse(grasp_wid_mask_target, inv_mat, w, h) for gi in range(len(num_grasps)): num_g = num_grasps[gi] grasp_preds, _ = detect_grasps( grasp_qua_mask_pred, grasp_sin_mask_pred, grasp_cos_mask_pred, grasp_wid_mask_pred, num_g ) j_index = calculate_jacquard_index(grasp_preds, grasp_target) num_correct_grasps[gi] += j_index num_total_grasps[gi] += 1 current_j1 = 100.0 * num_correct_grasps[0] / max(num_total_grasps[0], 1) current_j5 = 100.0 * num_correct_grasps[1] / max(num_total_grasps[1], 1) pbar.set_description( f"Evaluation: Epoch=[{epoch}/{args.epochs}] J@1: {current_j1:.2f} J@5: {current_j5:.2f}" ) J_index = [0, 0] for i in range(len(num_grasps)): J_index[i] = num_correct_grasps[i] / max(num_total_grasps[i], 1) head = 'Evaluation: Epoch=[{}/{}] J_index@1: {:.2f} J_index@5: {:.2f}'.format( epoch, args.epochs, 100. * J_index[0], 100. * J_index[1] ) logger.info(head) return 0.0, {}, J_index @torch.no_grad() def validate_without_grasp(val_loader, model, epoch, args): model.eval() time.sleep(2) pbar = tqdm(val_loader) for data in pbar: image = data["img"] text = data["word_vec"] grasp_qua_mask = data["grasp_masks"]["qua"] grasp_sin_mask = data["grasp_masks"]["sin"] grasp_cos_mask = data["grasp_masks"]["cos"] grasp_wid_mask = data["grasp_masks"]["wid"] image = image.cuda(non_blocking=True) text = text.cuda(non_blocking=True) grasp_qua_mask = grasp_qua_mask.cuda(non_blocking=True).unsqueeze(1) grasp_sin_mask = grasp_sin_mask.cuda(non_blocking=True).unsqueeze(1) grasp_cos_mask = grasp_cos_mask.cuda(non_blocking=True).unsqueeze(1) grasp_wid_mask = grasp_wid_mask.cuda(non_blocking=True).unsqueeze(1) _pred, _target = model(image, text, grasp_qua_mask, grasp_sin_mask, grasp_cos_mask, grasp_wid_mask) head = 'Evaluation: Epoch=[{}/{}] (no ins mask; validate_without_grasp does not compute metrics)'.format( epoch, args.epochs ) logger.info(head) return 0.0, {}, [0.0, 0.0] @torch.no_grad() def inference_with_grasp(test_loader, model, args): def inverse(img, mat, w, h): img = np.asarray(img) if img.dtype != np.float32: img = img.astype(np.float32) if torch.is_tensor(mat): mat = mat.detach().cpu().numpy() mat = np.asarray(mat, dtype=np.float32) if torch.is_tensor(w): w = int(w.item()) else: w = int(w) if torch.is_tensor(h): h = int(h.item()) else: h = int(h) inv_img = cv2.warpAffine( img, mat, (w, h), flags=cv2.INTER_CUBIC, borderValue=0. ) return inv_img model.eval() time.sleep(2) num_grasps = [1, 5] num_correct_grasps = [0, 0] num_total_grasps = [0, 0] iou_list = [] # keep for minimal diff tbar = tqdm(test_loader, desc='Inference:', ncols=100) for cnt, data in enumerate(tbar): image = data["img"] # ✅ 对齐 train/val：支持 raw text text = data["word_vec"] if not getattr(args, 'use_raw_text', False) else data["sentence"] grasp_qua_mask = data["grasp_masks"]["qua"] grasp_sin_mask = data["grasp_masks"]["sin"] grasp_cos_mask = data["grasp_masks"]["cos"] grasp_wid_mask = data["grasp_masks"]["wid"] inverse_matrix = data["inverse"] ori_sizes = data["ori_size"] grasp_targets = data["grasps"] sentences = data.get("sentence", None) img_paths = data.get("img_path", None) image = image.cuda(non_blocking=True) if isinstance(text, torch.Tensor): text = text.cuda(non_blocking=True) grasp_qua_mask = grasp_qua_mask.cuda(non_blocking=True).unsqueeze(1) grasp_sin_mask = grasp_sin_mask.cuda(non_blocking=True).unsqueeze(1) grasp_cos_mask = grasp_cos_mask.cuda(non_blocking=True).unsqueeze(1) grasp_wid_mask = grasp_wid_mask.cuda(non_blocking=True).unsqueeze(1) with amp.autocast(): pred, target = model( image, text, grasp_qua_mask, grasp_sin_mask, grasp_cos_mask, grasp_wid_mask ) grasp_qua_mask_preds = torch.sigmoid(pred[0]) grasp_sin_mask_preds = pred[1] grasp_cos_mask_preds = pred[2] grasp_wid_mask_preds = torch.sigmoid(pred[3]) # resize preds to input image size if grasp_qua_mask_preds.shape[-2:] != image.shape[-2:]: grasp_qua_mask_preds = F.interpolate( grasp_qua_mask_preds, size=image.shape[-2:], mode='bicubic', align_corners=True ).squeeze(1) grasp_sin_mask_preds = F.interpolate( grasp_sin_mask_preds, size=image.shape[-2:], mode='bicubic', align_corners=True ).squeeze(1) grasp_cos_mask_preds = F.interpolate( grasp_cos_mask_preds, size=image.shape[-2:], mode='bicubic', align_corners=True ).squeeze(1) grasp_wid_mask_preds = F.interpolate( grasp_wid_mask_preds, size=image.shape[-2:], mode='bicubic', align_corners=True ).squeeze(1) else: grasp_qua_mask_preds = grasp_qua_mask_preds.squeeze(1) grasp_sin_mask_preds = grasp_sin_mask_preds.squeeze(1) grasp_cos_mask_preds = grasp_cos_mask_preds.squeeze(1) grasp_wid_mask_preds = grasp_wid_mask_preds.squeeze(1) for idx in range(grasp_qua_mask_preds.shape[0]): inv_mat = inverse_matrix[idx] h, w = ori_sizes[idx] sent = sentences[idx] if sentences is not None else "" img_path = img_paths[idx] if img_paths is not None else None grasp_qua_mask_pred = grasp_qua_mask_preds[idx].cpu().numpy() grasp_sin_mask_pred = grasp_sin_mask_preds[idx].cpu().numpy() grasp_cos_mask_pred = grasp_cos_mask_preds[idx].cpu().numpy() grasp_wid_mask_pred = grasp_wid_mask_preds[idx].cpu().numpy() grasp_target = grasp_targets[idx] grasp_qua_mask_pred = inverse(grasp_qua_mask_pred, inv_mat, w, h) grasp_sin_mask_pred = inverse(grasp_sin_mask_pred, inv_mat, w, h) grasp_cos_mask_pred = inverse(grasp_cos_mask_pred, inv_mat, w, h) grasp_wid_mask_pred = inverse(grasp_wid_mask_pred, inv_mat, w, h) for gi, num_g in enumerate(num_grasps): grasp_preds, grasp_ang_mask_pred = detect_grasps( grasp_qua_mask_pred, grasp_sin_mask_pred, grasp_cos_mask_pred, grasp_wid_mask_pred, num_g ) j_index = calculate_jacquard_index(grasp_preds, grasp_target) num_correct_grasps[gi] += j_index num_total_grasps[gi] += 1 if getattr(args, "visualize", False) and img_path is not None: img_bgr = cv2.imread(img_path) img = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB) visualization( img, None, # ins_mask_pred removed (grasp_qua_mask_pred, grasp_ang_mask_pred, grasp_wid_mask_pred), grasp_preds, sent, save_path=os.path.join("./results", args.exp_name, f"results_{cnt}_{num_g}_grasps.png") ) J_index = [0.0, 0.0] for i in range(len(num_grasps)): J_index[i] = num_correct_grasps[i] / max(num_total_grasps[i], 1) logger.info("J@1: {:.2f}, J@5: {:.2f}".format(100. * J_index[0], 100. * J_index[1])) return 0.0, {}, J_index