| import warnings | |
| import pytorch_lightning as pl | |
| from realfake.config import SEED | |
| from realfake.models import RealFakeClassifier, RealFakeDataModule, RealFakeParams | |
| from realfake.train import prepare_trainer | |
| def main() -> None: | |
| pl.seed_everything(SEED) | |
| args = RealFakeParams.from_args() | |
| model = RealFakeClassifier(args) | |
| data = RealFakeDataModule(args.jsonl_file, args.batch_size, args.accelerator.devices * 4) | |
| trainer = prepare_trainer(args) | |
| if args.dry_run: | |
| print("Dry run, skipping training.") | |
| print("Model summary:") | |
| print(model) | |
| print("Data summary:") | |
| data.setup() | |
| print("Train batches:", len(data.dls[0])) | |
| print("Valid batches:", len(data.dls[1])) | |
| else: | |
| trainer.fit(model, datamodule=data) | |
| if __name__ == "__main__": | |
| with warnings.catch_warnings(): | |
| warnings.filterwarnings("ignore", category=Warning) | |
| main() | |