import argparse import json from pathlib import Path import torch from PIL import Image from torch.utils.data import Dataset, DataLoader from tqdm import tqdm from transformers import RTDetrImageProcessor, RTDetrForObjectDetection BASE_MODEL = "PekingU/rtdetr_r50vd" def load_classes(path): return [x.strip() for x in Path(path).read_text().splitlines() if x.strip()] class COCODetectionDataset(Dataset): def __init__(self, image_dir, annotation_file, processor): self.image_dir = Path(image_dir) self.processor = processor coco = json.loads(Path(annotation_file).read_text()) self.images = {x["id"]: x for x in coco["images"]} cats = sorted(coco["categories"], key=lambda x: x["id"]) self.category_id_to_label = {c["id"]: i for i,c in enumerate(cats)} anns = {} for a in coco["annotations"]: if not a.get("iscrowd", 0): anns.setdefault(a["image_id"], []).append(a) self.records = [] for image_id, info in self.images.items(): self.records.append({ "image_id": image_id, "file_name": info["file_name"], "width": info["width"], "height": info["height"], "annotations": anns.get(image_id, []) }) def __len__(self): return len(self.records) def __getitem__(self, idx): r = self.records[idx] image = Image.open(self.image_dir / r["file_name"]).convert("RGB") anns = [] for a in r["annotations"]: x,y,w,h = a["bbox"] if w <= 0 or h <= 0: continue anns.append({ "id": a["id"], "image_id": int(idx), "category_id": self.category_id_to_label[a["category_id"]], "bbox": [x,y,w,h], "area": float(a.get("area",w*h)), "iscrowd": 0 }) encoded = self.processor( images=image, annotations={"image_id": int(idx), "annotations": anns}, return_tensors="pt" ) encoded["pixel_values"] = encoded["pixel_values"].squeeze(0) if "pixel_mask" in encoded: encoded["pixel_mask"] = encoded["pixel_mask"].squeeze(0) encoded["labels"] = encoded["labels"][0] return encoded def collate_fn(batch): out = {"pixel_values": torch.stack([x["pixel_values"] for x in batch]), "labels": [x["labels"] for x in batch]} if "pixel_mask" in batch[0]: out["pixel_mask"] = torch.stack([x["pixel_mask"] for x in batch]) return out def move_to_device(obj, device): if torch.is_tensor(obj): return obj.to(device) if isinstance(obj, dict): return {k: move_to_device(v, device) for k, v in obj.items()} if isinstance(obj, list): return [move_to_device(v, device) for v in obj] if isinstance(obj, tuple): return tuple(move_to_device(v, device) for v in obj) return obj def evaluate(model, loader, device): model.eval(); total=0; n=0 with torch.no_grad(): for batch in loader: batch=move_to_device(batch, device) total += float(model(**batch).loss.item()); n += 1 model.train() return total/max(n,1) def patch_rtdetr_denoising_device(): """Patch RT-DETR denoising so embedding indices are always on the embedding device. Important: the wrapper must be an nn.Module, not a plain Python function. Transformers may inspect/use class_embed as a module, and a plain function loses the original module parameters/device information. """ try: import transformers.models.rt_detr.modeling_rt_detr as rtdetr_mod import torch.nn as nn original = rtdetr_mod.get_contrastive_denoising_training_group if getattr(original, "_icecream_device_patch", False): return class DeviceSafeEmbedding(nn.Module): def __init__(self, embedding): super().__init__() self.embedding = embedding def forward(self, indices): if torch.is_tensor(indices): indices = indices.to(self.embedding.weight.device) return self.embedding(indices) def wrapped(targets, num_classes, num_queries, class_embed, num_denoising_queries=100, label_noise_ratio=0.5, box_noise_scale=1.0, **kwargs): # Move every target tensor to the embedding/model device. try: embed_device = class_embed.weight.device except Exception: try: embed_device = next(class_embed.parameters()).device except Exception: embed_device = None if embed_device is not None: fixed_targets = [] for target in targets: if isinstance(target, dict): target = dict(target) for key, value in list(target.items()): if torch.is_tensor(value): target[key] = value.to(embed_device) fixed_targets.append(target) targets = fixed_targets class_embed = DeviceSafeEmbedding(class_embed) return original( targets=targets, num_classes=num_classes, num_queries=num_queries, class_embed=class_embed, num_denoising_queries=num_denoising_queries, label_noise_ratio=label_noise_ratio, box_noise_scale=box_noise_scale, **kwargs, ) wrapped._icecream_device_patch = True rtdetr_mod.get_contrastive_denoising_training_group = wrapped print("RT-DETR denoising device patch installed") except Exception as exc: raise RuntimeError(f"Could not install RT-DETR denoising device patch: {exc}") from exc def main(): p=argparse.ArgumentParser() p.add_argument("--train-dir",required=True); p.add_argument("--val-dir",required=True) p.add_argument("--classes",required=True); p.add_argument("--output-dir",default="model") p.add_argument("--epochs",type=int,default=30); p.add_argument("--batch-size",type=int,default=2) p.add_argument("--learning-rate",type=float,default=1e-5); p.add_argument("--weight-decay",type=float,default=1e-4) p.add_argument("--num-workers",type=int,default=2) a=p.parse_args() classes=load_classes(a.classes) id2label={i:n for i,n in enumerate(classes)} label2id={n:i for i,n in enumerate(classes)} proc=RTDetrImageProcessor.from_pretrained(BASE_MODEL) train=COCODetectionDataset(Path(a.train_dir)/"images",Path(a.train_dir)/"annotations.json",proc) val=COCODetectionDataset(Path(a.val_dir)/"images",Path(a.val_dir)/"annotations.json",proc) if len(train)==0 or len(val)==0: raise ValueError("Training and validation datasets must contain at least one image.") if len(train.category_id_to_label)!=len(classes) or len(val.category_id_to_label)!=len(classes): raise ValueError("COCO categories do not match classes.txt. Rebuild the dataset after saving the classes.") model=RTDetrForObjectDetection.from_pretrained( BASE_MODEL,num_labels=len(classes),id2label=id2label,label2id=label2id, ignore_mismatched_sizes=True ) # Some Transformers RT-DETR releases enter the denoising path whenever # training, regardless of num_denoising. Keep the config disabled where # supported, but also install the device-safe embedding patch below. for cfg_owner in (model, getattr(model, "model", None)): cfg = getattr(cfg_owner, "config", None) if cfg is not None: for name in ("num_denoising", "num_denoising_queries"): if hasattr(cfg, name): setattr(cfg, name, 0) device=torch.device("cuda" if torch.cuda.is_available() else "cpu") model.to(device) patch_rtdetr_denoising_device() tr=DataLoader(train,batch_size=a.batch_size,shuffle=True,num_workers=0,collate_fn=collate_fn) va=DataLoader(val,batch_size=a.batch_size,shuffle=False,num_workers=0,collate_fn=collate_fn) opt=torch.optim.AdamW(model.parameters(),lr=a.learning_rate,weight_decay=a.weight_decay) outdir=Path(a.output_dir); outdir.mkdir(parents=True,exist_ok=True) best=float("inf") for epoch in range(a.epochs): model.train(); running=0 bar=tqdm(tr,desc=f"epoch {epoch+1}/{a.epochs}") for step,batch in enumerate(bar): batch=move_to_device(batch, device) # RT-DETR's loss matcher uses nested target tensors (boxes/classes). # Move every tensor in labels to the same device as the model. if "labels" in batch: # RT-DETR expects every nested target tensor on the same device as the model. for target in batch["labels"]: if isinstance(target, dict): for key, value in list(target.items()): if torch.is_tensor(value): target[key] = value.to(device) loss=model(**batch).loss loss.backward(); opt.step(); opt.zero_grad(set_to_none=True) running += float(loss.item()) bar.set_postfix(loss=f"{running/(step+1):.4f}") vl=evaluate(model,va,device) print(f"validation_loss={vl:.4f}") if vl