import os import time import torch import importlib from loguru import logger from utils import seed_torch, evaluate from dataSets import TrainDataset import warnings warnings.filterwarnings("ignore") class Trainer: def __init__(self, mode: str, device: str, branch: str, train_dataset: str, label_smooth: bool = False, ckpt: str = "ckpt", epochs: int = 20, batch_size: int = 32): self.device = device self.mode = mode self.branch = branch self.train_dataset = train_dataset self.label_smooth = label_smooth self.ckpt = ckpt self.epochs = epochs self.batch_size = batch_size # Dynamically import the detector module from either "progan" or "sd-v1_4" detector_module = importlib.import_module(f"detectors.{train_dataset}") # Get the detector based on mode if self.mode == "branch": ArtifactDetector = getattr(detector_module, "ArtifactDetector") SemanticDetector = getattr(detector_module, "SemanticDetector") if self.branch == "artifact": self.model = ArtifactDetector() elif self.branch == "semantic": self.model = SemanticDetector() else: raise ValueError(f"Unknown detector: {self.branch}") elif self.mode == "fusion": semantic_weights_path = os.path.join(self.ckpt, self.train_dataset, "semantic", "best_model.pth") artifact_weights_path = os.path.join(self.ckpt, self.train_dataset, "artifact", "best_model.pth") if not os.path.exists(semantic_weights_path) or not os.path.exists(artifact_weights_path): raise ValueError("Semantic or Artifact weights path does not exist for fusion mode") CoSpyFusionDetector = getattr(detector_module, "CoSpyFusionDetector") self.model = CoSpyFusionDetector( semantic_weights_path=semantic_weights_path, artifact_weights_path=artifact_weights_path) elif self.mode == "end2end": End2EndDetector = getattr(detector_module, "End2EndDetector") self.model = End2EndDetector() else: raise ValueError(f"Unknown mode: {self.mode}") self.model.to(self.device) # Initialize the fc layer torch.nn.init.normal_(self.model.fc.weight.data, 0.0, 0.02) if self.mode == "end2end": torch.nn.init.normal_(self.model.sem.fc.weight.data, 0.0, 0.02) torch.nn.init.normal_(self.model.art.fc.weight.data, 0.0, 0.02) # Optimizer _beta1 = 0.9 _weight_decay = 0.0 params = [p for p in self.model.parameters() if p.requires_grad] logger.info(f"Trainable parameters: {len(params)}") self._lr = 1e-4 if self.mode != "fusion" else 1e-1 self.optimizer = torch.optim.AdamW(params, lr=self._lr, betas=(_beta1, 0.999), weight_decay=_weight_decay) # Loss function if self.label_smooth: self.criterion = LabelSmoothingBCEWithLogits(smoothing=0.1) else: self.criterion = torch.nn.BCEWithLogitsLoss() # Scheduler self.delr_freq = 10 def train_step(self, batch_data): inputs, labels = batch_data inputs, labels = inputs.to(self.device), labels.to(self.device) self.optimizer.zero_grad() outputs = self.model(inputs) loss = self.criterion(outputs, labels.unsqueeze(1).float()) loss.backward() self.optimizer.step() eval_loss = loss.item() y_pred = outputs.sigmoid().flatten().tolist() y_true = labels.tolist() return eval_loss, y_pred, y_true def scheduler(self, status_dict): epoch = status_dict["epoch"] if epoch % self.delr_freq == 0 and epoch != 0: for param_group in self.optimizer.param_groups: param_group["lr"] *= 0.9 self._lr = param_group["lr"] def train(self): # Determine data split and transform based on mode if self.mode == "fusion": train_split, val_split = "val", "val" train_transform = self.model.test_transform test_transform = self.model.test_transform else: # branch or end2end mode train_split, val_split = "train", "val" train_transform = self.model.train_transform test_transform = self.model.test_transform # Determine save directory if self.mode == "branch": subdir = self.branch else: subdir = self.mode # Set the saving directory model_dir = os.path.join(self.ckpt, self.train_dataset, subdir) if not os.path.exists(model_dir): os.makedirs(model_dir) # Setup logger log_path = f"{model_dir}/training.log" if os.path.exists(log_path): os.remove(log_path) logger_id = logger.add( log_path, format="{time:MM-DD at HH:mm:ss} | {level} | {module}:{line} | {message}", level="DEBUG", ) # Add JPEG compression for sd-v1_4 dataset self.add_jpeg = True if self.train_dataset == "sd-v1_4" else False # Load the training and validation dataset train_dataset = TrainDataset(train_dataset=self.train_dataset, split=train_split, add_jpeg=self.add_jpeg, transform=train_transform) train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=self.batch_size, shuffle=True, num_workers=4, pin_memory=True) val_dataset = TrainDataset(train_dataset=self.train_dataset, split=val_split, add_jpeg=self.add_jpeg, transform=test_transform) val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=self.batch_size, shuffle=False, num_workers=4, pin_memory=True) logger.info(f"Train size {len(train_dataset)} | Val size {len(val_dataset)}") # Train the detector best_acc = 0 for epoch in range(self.epochs): self.model.train() time_start = time.time() for step_id, batch_data in enumerate(train_loader): eval_loss, y_pred, y_true = self.train_step(batch_data) ap, accuracy = evaluate(y_pred, y_true) if (step_id + 1) % 100 == 0: time_end = time.time() logger.info(f"Epoch {epoch} | Batch {step_id + 1}/{len(train_loader)} | Loss {eval_loss:.4f} | AP {ap*100:.2f}% | Accuracy {accuracy*100:.2f}% | Time {time_end-time_start:.2f}s") time_start = time.time() # Evaluate the model self.model.eval() y_pred, y_true = [], [] for (images, labels) in val_loader: y_pred.extend(self.model.predict(images)) y_true.extend(labels.tolist()) ap, accuracy = evaluate(y_pred, y_true) eval_type = "Test" if self.mode == "branch" else "Total" logger.info(f"Epoch {epoch} | {eval_type} AP {ap*100:.2f}% | {eval_type} Accuracy {accuracy*100:.2f}%") # Schedule the training status_dict = {"epoch": epoch, "AP": ap, "Accuracy": accuracy} self.scheduler(status_dict) # Save the model if accuracy >= best_acc: best_acc = accuracy self.model.save_weights(f"{model_dir}/best_model.pth") logger.info(f"Best model saved with accuracy {best_acc*100:.2f}%") if epoch % 5 == 0: self.model.save_weights(f"{model_dir}/epoch_{epoch}.pth") logger.info(f"Model saved at epoch {epoch}") # Save the final model self.model.save_weights(f"{model_dir}/final_model.pth") logger.info("Final model saved") # Remove the logger logger.remove(logger_id)