add xception example

ControlNet · ControlNet · commit b4ce1b1c1a62 · 2025-04-24T03:00:49.000+10:00
diff --git a/examples/xception/README.md b/examples/xception/README.md
@@ -0,0 +1,23 @@
+# Xception
+
+This example trains a Xception model on the AVDeepfake1M/AVDeepfake1M++ dataset for classification with video-level labels.
+
+## Requirements
+
+- Python
+- PyTorch
+- PyTorch Lightning
+- TIMM
+- AVDeepfake1M SDK
+
+
+## Training
+
+```bash
+python train.py --data_root /path/to/avdeepfake1m --model xception
+```
+
+## Output
+
+*   **Checkpoints:** Model checkpoints are saved under `./ckpt1/xception/`. The last checkpoint is saved as `last.ckpt`.
+*   **Logs:** Training logs (including metrics like `train_loss`, `val_loss`, and learning rates) are saved by PyTorch Lightning, typically in a directory named `./lightning_logs/`. You can view these logs using TensorBoard (`tensorboard --logdir ./lightning_logs`). 
diff --git a/examples/xception/train.py b/examples/xception/train.py
@@ -0,0 +1,94 @@
+import argparse
+
+from torch.utils.data import DataLoader
+from pytorch_lightning import Trainer
+from pytorch_lightning.callbacks import ModelCheckpoint
+from avdeepfake1m.loader import AVDeepfake1mPlusPlusImages
+
+from xception import Xception
+from utils import LrLogger, EarlyStoppingLR
+
+
+parser = argparse.ArgumentParser(description="Classification model training")
+parser.add_argument("--data_root", type=str)
+parser.add_argument("--batch_size", type=int, default=128)
+parser.add_argument("--model", type=str, choices=["xception", "meso4", "meso_inception4"])
+parser.add_argument("--gpus", type=int, default=1)
+parser.add_argument("--precision", default=32)
+parser.add_argument("--num_train", type=int, default=None)
+parser.add_argument("--num_val", type=int, default=2000)
+parser.add_argument("--max_epochs", type=int, default=500)
+parser.add_argument("--resume", type=str, default=None)
+args = parser.parse_args()
+
+
+if __name__ == "__main__":
+
+    # You can fix the random seed if you want reproducible subsets each epoch:
+    # torch.manual_seed(42)
+    # random.seed(42)
+
+    learning_rate = 1e-4
+    gpus = args.gpus
+    total_batch_size = args.batch_size * gpus
+    learning_rate = learning_rate * total_batch_size / 4
+
+    # Setup model
+    if args.model == "xception":
+        model = Xception(learning_rate, distributed=gpus > 1)
+    else:
+        raise ValueError(f"Unknown model: {args.model}")
+
+    train_dataset = AVDeepfake1mPlusPlusImages(
+        subset="train",
+        data_root=args.data_root,
+        take_num=args.num_train,
+        use_video_label=True  # For video-level label access, set True
+    )
+
+    # For validation, you can still do the normal dataset
+    val_dataset = AVDeepfake1mPlusPlusImages(
+        subset="val",
+        data_root=args.data_root,
+        take_num=args.num_val,
+        use_video_label=True
+    )
+
+    # Parse precision properly
+    try:
+        precision = int(args.precision)
+    except ValueError:
+        precision = args.precision
+
+    monitor = "val_loss"
+
+    trainer = Trainer(
+        log_every_n_steps=50,
+        precision=precision,
+        max_epochs=args.max_epochs,
+        callbacks=[
+            ModelCheckpoint(
+                dirpath=f"./ckpt1/{args.model}",
+                save_last=True,
+                filename=args.model + "-{epoch}-{val_loss:.3f}",
+                monitor=monitor,
+                mode="min"
+            ),
+            LrLogger(),
+            EarlyStoppingLR(lr_threshold=1e-7)
+        ],
+        enable_checkpointing=True,
+        benchmark=True,
+        accelerator="gpu",
+        devices=args.gpus,
+        strategy="ddp" if args.gpus > 1 else "auto",
+        # ckpt_path=args.resume,
+        # If you're on an older version of Lightning, you may need `strategy='ddp'` just the same, but this is typical.
+    )
+
+    trainer.fit(
+        model,
+        train_dataloaders=DataLoader(train_dataset, batch_size=args.batch_size, num_workers=0),
+        val_dataloaders=DataLoader(val_dataset, batch_size=args.batch_size, num_workers=0),
+        ckpt_path=args.resume,
+    )
diff --git a/examples/xception/utils.py b/examples/xception/utils.py
@@ -0,0 +1,49 @@
+import re
+
+from pytorch_lightning import Callback, Trainer, LightningModule
+
+
+class LrLogger(Callback):
+    """Log learning rate in each epoch start."""
+
+    def on_train_epoch_start(self, trainer: Trainer, pl_module: LightningModule) -> None:
+        for i, optimizer in enumerate(trainer.optimizers):
+            for j, params in enumerate(optimizer.param_groups):
+                key = f"opt{i}_lr{j}"
+                value = params["lr"]
+                pl_module.logger.log_metrics({key: value}, step=trainer.global_step)
+                pl_module.log(key, value, logger=False, sync_dist=pl_module.distributed)
+
+
+class EarlyStoppingLR(Callback):
+    """Early stop model training when the LR is lower than threshold."""
+
+    def __init__(self, lr_threshold: float, mode="all"):
+        self.lr_threshold = lr_threshold
+
+        if mode in ("any", "all"):
+            self.mode = mode
+        else:
+            raise ValueError(f"mode must be one of ('any', 'all')")
+
+    def on_train_epoch_start(self, trainer: Trainer, pl_module: LightningModule) -> None:
+        self._run_early_stop_checking(trainer)
+
+    def _run_early_stop_checking(self, trainer: Trainer) -> None:
+        metrics = trainer._logger_connector.callback_metrics
+        if len(metrics) == 0:
+            return
+        all_lr = []
+        for key, value in metrics.items():
+            if re.match(r"opt\d+_lr\d+", key):
+                all_lr.append(value)
+
+        if len(all_lr) == 0:
+            return
+
+        if self.mode == "all":
+            if all(lr <= self.lr_threshold for lr in all_lr):
+                trainer.should_stop = True
+        elif self.mode == "any":
+            if any(lr <= self.lr_threshold for lr in all_lr):
+                trainer.should_stop = True
diff --git a/examples/xception/xception.py b/examples/xception/xception.py
@@ -0,0 +1,36 @@
+import timm
+
+from pytorch_lightning import LightningModule
+from torch.nn import BCEWithLogitsLoss
+from torch.optim import Adam
+
+
+class Xception(LightningModule):
+    def __init__(self, lr, distributed=False):
+        super(Xception, self).__init__()
+        self.lr = lr
+        self.model = timm.create_model('xception', pretrained=True, num_classes=1)
+        self.loss_fn = BCEWithLogitsLoss()
+        self.distributed = distributed
+
+    def forward(self, x):
+        x = self.model(x)
+        return x
+
+    def training_step(self, batch, batch_idx):
+        x, y = batch
+        y_hat = self(x)
+        loss = self.loss_fn(y_hat, y.unsqueeze(1))
+        self.log('train_loss', loss)
+        return loss
+
+    def validation_step(self, batch, batch_idx):
+        x, y = batch
+        y_hat = self(x)
+        loss = self.loss_fn(y_hat, y.unsqueeze(1))
+        self.log('val_loss', loss)
+        return loss
+
+    def configure_optimizers(self):
+        optimizer = Adam(self.parameters(), lr=self.lr)
+        return [optimizer]
diff --git a/python/avdeepfake1m/loader.py b/python/avdeepfake1m/loader.py
@@ -538,3 +538,64 @@ def sample_indexes(total_frames: int, n_frames: int, temporal_sample_rate: int):
         print(f"total_frames: {total_frames}, n_frames: {n_frames}, temporal_sample_rate: {temporal_sample_rate}")
         raise e
     return torch.arange(n_frames) * temporal_sample_rate + start_ind
+
+
+class AVDeepfake1mPlusPlusImages(IterableDataset):
+
+    def __init__(self, subset: str, data_root: str = "data",
+        image_size: int = 96,
+        use_video_label: bool = False,
+        use_seg_label: Optional[int] = None,
+        take_num: Optional[int] = None,
+        metadata: Optional[List[Metadata]] = None,
+    ):
+        self.subset = subset
+        self.data_root = data_root
+        self.image_size = image_size
+        self.use_video_label = use_video_label
+        if self.use_video_label:
+            assert use_seg_label is None
+        self.use_seg_label = use_seg_label
+        if metadata is None:
+            metadata_json = read_json(os.path.join(self.data_root, f"{subset}_metadata.json"))
+            self.metadata = [Metadata(**meta, fps=25) for meta in metadata_json]
+        else:
+            self.metadata = metadata
+
+        if take_num is not None:
+            self.metadata = self.metadata[:take_num]
+
+        self.total_frames = sum([each.video_frames for each in self.metadata])
+        print("Load {} data in {}.".format(len(self.metadata), subset))
+
+    def __len__(self):
+        return self.total_frames
+
+    def __iter__(self):
+        for meta in self.metadata:
+            video = read_video_fast(os.path.join(self.data_root,self.subset, meta.file))
+            if self.image_size != 224:
+                video = resize_video(video, (96, 96))
+            if self.use_video_label:
+                label = float(len(meta.fake_periods) > 0)
+                for frame in video:
+                    yield frame, label
+            elif self.use_seg_label:
+                frame_label = torch.zeros(len(video))
+                for begin, end in meta.fake_periods:
+                    begin = int(begin * 25)
+                    end = int(end * 25)
+                    frame_label[begin: end] = 1
+                seg_label = torch.split(frame_label, self.use_seg_label)
+                seg_label = torch.nn.utils.rnn.pad_sequence(seg_label, batch_first=True)
+                seg_label = (seg_label.sum(dim=1) > 0).float().repeat_interleave(self.use_seg_label)
+                for i, frame in enumerate(video):
+                    yield frame, seg_label[i]
+            else:
+                frame_label = torch.zeros(len(video))
+                for begin, end in meta.fake_periods:
+                    begin = int(begin * 25)
+                    end = int(end * 25)
+                    frame_label[begin: end] = 1
+                for i, frame in enumerate(video):
+                    yield frame, frame_label[i]
