54 lines
1.7 KiB
Python
54 lines
1.7 KiB
Python
#!/usr/bin/env python
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
from pathlib import Path
|
|
|
|
from torch.utils.data import DataLoader, random_split
|
|
|
|
from flux2_resolution_guard.config import TrainConfig
|
|
from flux2_resolution_guard.data import SyntheticFlux2DriftDataset
|
|
from flux2_resolution_guard.models import SMICConfig, SMICCorrectionModel
|
|
from flux2_resolution_guard.training import Trainer
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--image-dir", required=True)
|
|
parser.add_argument("--output-dir", required=True)
|
|
parser.add_argument("--image-size", type=int, default=256)
|
|
parser.add_argument("--epochs", type=int, default=10)
|
|
parser.add_argument("--batch-size", type=int, default=4)
|
|
parser.add_argument("--device", default="cuda")
|
|
args = parser.parse_args()
|
|
|
|
dataset = SyntheticFlux2DriftDataset(
|
|
image_dir=args.image_dir,
|
|
image_size=args.image_size,
|
|
)
|
|
n_val = max(1, int(len(dataset) * 0.1))
|
|
n_train = max(1, len(dataset) - n_val)
|
|
train_ds, val_ds = random_split(dataset, [n_train, n_val])
|
|
|
|
train_loader = DataLoader(train_ds, batch_size=args.batch_size, shuffle=True, num_workers=0)
|
|
val_loader = DataLoader(val_ds, batch_size=args.batch_size, shuffle=False, num_workers=0)
|
|
|
|
model = SMICCorrectionModel(SMICConfig())
|
|
trainer = Trainer(
|
|
model=model,
|
|
train_loader=train_loader,
|
|
val_loader=val_loader,
|
|
output_dir=args.output_dir,
|
|
config=TrainConfig(
|
|
image_size=args.image_size,
|
|
epochs=args.epochs,
|
|
batch_size=args.batch_size,
|
|
device=args.device,
|
|
),
|
|
)
|
|
trainer.fit()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|