File size: 789 Bytes
31cab2b 649c592 b80dc1e 31cab2b b80dc1e d5ba31a 649c592 1ff26a6 b80dc1e 5fed6e3 1ff26a6 649c592 5fed6e3 649c592 b80dc1e 642e32f 5fed6e3 |
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 |
import hydra
import torch
from loguru import logger
from config.config import Config
from model.yolo import get_model
from tools.log_helper import custom_logger
from tools.trainer import Trainer
from utils.dataloader import get_dataloader
from utils.get_dataset import prepare_dataset
@hydra.main(config_path="config", config_name="config", version_base=None)
def main(cfg: Config):
if cfg.download.auto:
prepare_dataset(cfg.download)
dataloader = get_dataloader(cfg)
model = get_model(cfg.model)
# TODO: get_device or rank, for DDP mode
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
trainer = Trainer(model, cfg.hyper.train, device)
trainer.train(dataloader, 10)
if __name__ == "__main__":
custom_logger()
main()
|