| """ |
| main.py |
| ======== |
| اسکریپت اصلی برای اجرای کل pipeline از خط فرمان: |
| python scripts/main.py --mode train |
| python scripts/main.py --mode evaluate --ckpt models/saved_models/omnial_epoch50.pt |
| """ |
|
|
| import os |
| import sys |
| import argparse |
| import yaml |
| import torch |
|
|
| sys.path.append(os.path.join(os.path.dirname(__file__), "..")) |
| from data.data_loader import build_dataloader |
| from scripts.train import build_model_and_optimizer, train_model |
| from scripts.evaluate import evaluate_dataset |
|
|
|
|
| def load_config(path: str = "config/config.yaml") -> dict: |
| with open(path, "r", encoding="utf-8") as f: |
| return yaml.safe_load(f) |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--mode", choices=["train", "evaluate"], default="train") |
| parser.add_argument("--dataset", choices=["mvtec", "visa"], default="mvtec") |
| parser.add_argument("--config", default="config/config.yaml") |
| parser.add_argument("--ckpt", default=None) |
| args = parser.parse_args() |
|
|
| cfg = load_config(args.config) |
| torch.manual_seed(cfg["train"]["seed"]) |
|
|
| model, optimizer, criterion, device = build_model_and_optimizer(cfg) |
| print(f"Using device: {device}") |
|
|
| if args.mode == "train": |
| train_loader, train_dataset = build_dataloader(cfg, dataset_name=args.dataset, shuffle=True) |
| print(f"تعداد نمونههای آموزشی ({args.dataset}): {len(train_dataset)}") |
| history = train_model(model, train_loader, optimizer, criterion, device, cfg) |
| print("آموزش تمام شد.") |
|
|
| elif args.mode == "evaluate": |
| if args.ckpt: |
| model.load_state_dict(torch.load(args.ckpt, map_location=device)) |
| print(f"وزنها از {args.ckpt} بارگذاری شد.") |
| evaluate_dataset(model, cfg, dataset_name=args.dataset, device=device) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
