Upload train.py

training/train.py

@@ -8,273 +8,125 @@ 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):
-    classes = [x.strip() for x in Path(path).read_text().splitlines() if x.strip()]
-    if not classes:
-        raise ValueError("classes.txt is empty")
-    return classes
-
+    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"]}
-
-        # The training model needs contiguous labels 0..N-1.
-        # categories.json is converted here from the COCO category IDs.
-        categories = sorted(coco["categories"], key=lambda x: x["id"])
-        self.category_id_to_label = {
-            cat["id"]: i for i, cat in enumerate(categories)
-        }
-
-        anns_by_image = {}
-        for ann in coco["annotations"]:
-            if ann.get("iscrowd", 0):
-                continue
-            anns_by_image.setdefault(ann["image_id"], []).append(ann)
-
+        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_by_image.get(image_id, []),
-                }
-            )
+            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 __len__(self): return len(self.records)
 
     def __getitem__(self, idx):
-        record = self.records[idx]
-        image_path = self.image_dir / record["file_name"]
-        image = Image.open(image_path).convert("RGB")
-
-        annotations = []
-        for ann in record["annotations"]:
-            x, y, w, h = ann["bbox"]
-            if w <= 0 or h <= 0:
-                continue
-
-            label = self.category_id_to_label[ann["category_id"]]
-            annotations.append(
-                {
-                    "id": ann["id"],
-                    "image_id": record["image_id"],
-                    "category_id": label,
-                    "bbox": [x, y, w, h],
-                    "area": float(ann.get("area", w * h)),
-                    "iscrowd": 0,
-                }
-            )
-
-        target = {
-            "image_id": record["image_id"],
-            "annotations": annotations,
-        }
-
+        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": r["image_id"],
+                "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=target,
-            return_tensors="pt",
+            annotations={"image_id": r["image_id"], "annotations": anns},
+            return_tensors="pt"
         )
-
-        # Remove batch dimension. DataLoader will create it.
         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):
-    pixel_values = torch.stack([x["pixel_values"] for x in batch])
-    pixel_mask = None
+    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]:
-        pixel_mask = torch.stack([x["pixel_mask"] for x in batch])
-
-    labels = [x["labels"] for x in batch]
-    result = {"pixel_values": pixel_values, "labels": labels}
-    if pixel_mask is not None:
-        result["pixel_mask"] = pixel_mask
-    return result
-
+        out["pixel_mask"] = torch.stack([x["pixel_mask"] for x in batch])
+    return out
 
 def evaluate(model, loader, device):
-    model.eval()
-    total = 0.0
-    count = 0
+    model.eval(); total=0; n=0
     with torch.no_grad():
         for batch in loader:
-            batch = {
-                k: (v.to(device) if torch.is_tensor(v) else v)
-                for k, v in batch.items()
-            }
-            out = model(**batch)
-            total += float(out.loss.item())
-            count += 1
+            batch={k:(v.to(device) if torch.is_tensor(v) else v) for k,v in batch.items()}
+            total += float(model(**batch).loss.item()); n += 1
     model.train()
-    return total / max(count, 1)
-
+    return total/max(n,1)
 
 def main():
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--train-dir", default="data/train")
-    parser.add_argument("--val-dir", default="data/val")
-    parser.add_argument("--classes", default="training/classes.txt")
-    parser.add_argument("--output-dir", default="model")
-    parser.add_argument("--epochs", type=int, default=30)
-    parser.add_argument("--batch-size", type=int, default=2)
-    parser.add_argument("--learning-rate", type=float, default=1e-5)
-    parser.add_argument("--weight-decay", type=float, default=1e-4)
-    parser.add_argument("--num-workers", type=int, default=2)
-    parser.add_argument("--grad-accumulation", type=int, default=1)
-    parser.add_argument("--resume", default=None)
-    args = parser.parse_args()
-
-    classes = load_classes(args.classes)
-    label2id = {name: i for i, name in enumerate(classes)}
-    id2label = {i: name for i, name in enumerate(classes)}
-
-    train_ann = Path(args.train_dir) / "annotations.json"
-    val_ann = Path(args.val_dir) / "annotations.json"
-
-    processor = RTDetrImageProcessor.from_pretrained(BASE_MODEL)
-
-    train_ds = COCODetectionDataset(
-        Path(args.train_dir) / "images",
-        train_ann,
-        processor,
-    )
-    val_ds = COCODetectionDataset(
-        Path(args.val_dir) / "images",
-        val_ann,
-        processor,
-    )
-
-    # Make sure both datasets use exactly the class list supplied by the user.
-    train_categories = len(train_ds.category_id_to_label)
-    val_categories = len(val_ds.category_id_to_label)
-    if train_categories != len(classes) or val_categories != len(classes):
-        raise ValueError(
-            "The number of COCO categories does not match classes.txt. "
-            f"classes.txt={len(classes)}, train={train_categories}, val={val_categories}"
-        )
-
-    model = RTDetrForObjectDetection.from_pretrained(
-        BASE_MODEL,
-        num_labels=len(classes),
-        id2label=id2label,
-        label2id=label2id,
-        ignore_mismatched_sizes=True,
+    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("Both train and validation splits must contain annotated images.")
+    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.")
+
+    model=RTDetrForObjectDetection.from_pretrained(
+        BASE_MODEL,num_labels=len(classes),id2label=id2label,label2id=label2id,
+        ignore_mismatched_sizes=True
     )
-
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    device=torch.device("cuda" if torch.cuda.is_available() else "cpu")
     model.to(device)
 
-    train_loader = DataLoader(
-        train_ds,
-        batch_size=args.batch_size,
-        shuffle=True,
-        num_workers=args.num_workers,
-        collate_fn=collate_fn,
-        pin_memory=torch.cuda.is_available(),
-    )
-    val_loader = DataLoader(
-        val_ds,
-        batch_size=args.batch_size,
-        shuffle=False,
-        num_workers=args.num_workers,
-        collate_fn=collate_fn,
-        pin_memory=torch.cuda.is_available(),
-    )
-
-    optimizer = torch.optim.AdamW(
-        model.parameters(),
-        lr=args.learning_rate,
-        weight_decay=args.weight_decay,
-    )
-
-    scaler = torch.amp.GradScaler("cuda", enabled=torch.cuda.is_available())
-    start_epoch = 0
-
-    if args.resume:
-        checkpoint = torch.load(args.resume, map_location=device)
-        model.load_state_dict(checkpoint["model"])
-        optimizer.load_state_dict(checkpoint["optimizer"])
-        start_epoch = checkpoint["epoch"] + 1
-
-    output_dir = Path(args.output_dir)
-    output_dir.mkdir(parents=True, exist_ok=True)
-
-    best_val = float("inf")
-
-    for epoch in range(start_epoch, args.epochs):
-        model.train()
-        optimizer.zero_grad(set_to_none=True)
-        running = 0.0
-
-        pbar = tqdm(train_loader, desc=f"Epoch {epoch + 1}/{args.epochs}")
-        for step, batch in enumerate(pbar):
-            batch = {
-                k: (v.to(device) if torch.is_tensor(v) else v)
-                for k, v in batch.items()
-            }
-
-            with torch.amp.autocast(
-                device_type="cuda",
-                enabled=torch.cuda.is_available(),
-            ):
-                out = model(**batch)
-                loss = out.loss / args.grad_accumulation
-
-            scaler.scale(loss).backward()
-
-            if (step + 1) % args.grad_accumulation == 0:
-                scaler.step(optimizer)
-                scaler.update()
-                optimizer.zero_grad(set_to_none=True)
-
-            running += float(loss.item()) * args.grad_accumulation
-            pbar.set_postfix(loss=f"{running / (step + 1):.4f}")
-
-        val_loss = evaluate(model, val_loader, device)
-        print(f"epoch={epoch + 1} validation_loss={val_loss:.4f}")
-
-        checkpoint = {
-            "epoch": epoch,
-            "model": model.state_dict(),
-            "optimizer": optimizer.state_dict(),
-        }
-        torch.save(checkpoint, output_dir / "last_checkpoint.pt")
-
-        if val_loss < best_val:
-            best_val = val_loss
-            model.save_pretrained(output_dir)
-            processor.save_pretrained(output_dir)
-            (output_dir / "classes.json").write_text(
-                json.dumps(
-                    {"id2label": id2label, "label2id": label2id},
-                    indent=2,
-                )
-            )
-            print(f"Saved best model to {output_dir}")
-
-    # Always save the final model as well.
-    model.save_pretrained(output_dir)
-    processor.save_pretrained(output_dir)
-    print(f"Final model saved to {output_dir}")
-
-
-if __name__ == "__main__":
-    main()
+    tr=DataLoader(train,batch_size=a.batch_size,shuffle=True,num_workers=a.num_workers,collate_fn=collate_fn)
+    va=DataLoader(val,batch_size=a.batch_size,shuffle=False,num_workers=a.num_workers,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={k:(v.to(device) if torch.is_tensor(v) else v) for k,v in batch.items()}
+            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<best:
+            best=vl
+            model.save_pretrained(outdir)
+            proc.save_pretrained(outdir)
+            (outdir/"classes.json").write_text(json.dumps({"id2label":id2label,"label2id":label2id},indent=2))
+    model.save_pretrained(outdir); proc.save_pretrained(outdir)
+
+if __name__=="__main__": main()