File size: 987 Bytes
1197f7d f1bfd74 1197f7d f1bfd74 1197f7d f1bfd74 1197f7d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 |
import sys
from pathlib import Path
import hydra
project_root = Path(__file__).resolve().parent.parent
sys.path.append(str(project_root))
from yolo import (
Config,
ModelTrainer,
ProgressLogger,
create_converter,
create_dataloader,
create_model,
)
from yolo.utils.model_utils import get_device
@hydra.main(config_path="config", config_name="config", version_base=None)
def main(cfg: Config):
progress = ProgressLogger(cfg, exp_name=cfg.name)
device, use_ddp = get_device(cfg.device)
dataloader = create_dataloader(cfg.task.data, cfg.dataset, cfg.task.task, use_ddp)
model = create_model(cfg.model, class_num=cfg.dataset.class_num, weight_path=cfg.weight)
model = model.to(device)
converter = create_converter(cfg.model.name, model, cfg.model.anchor, cfg.image_size, device)
solver = ModelTrainer(cfg, model, converter, progress, device)
progress.start()
solver.solve(dataloader)
if __name__ == "__main__":
main()
|