henry000 commited on
Commit
de05146
·
1 Parent(s): 7854350

🔥 [Remove] train.py which outside of ./yolo

Browse files
Files changed (1) hide show
  1. train.py +0 -29
train.py DELETED
@@ -1,29 +0,0 @@
1
- import hydra
2
- import torch
3
- from loguru import logger
4
-
5
- from config.config import Config
6
- from model.yolo import get_model
7
- from tools.log_helper import custom_logger
8
- from tools.trainer import Trainer
9
- from utils.dataloader import get_dataloader
10
- from utils.get_dataset import prepare_dataset
11
-
12
-
13
- @hydra.main(config_path="config", config_name="config", version_base=None)
14
- def main(cfg: Config):
15
- if cfg.download.auto:
16
- prepare_dataset(cfg.download)
17
-
18
- dataloader = get_dataloader(cfg)
19
- model = get_model(cfg.model)
20
- # TODO: get_device or rank, for DDP mode
21
- device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
22
-
23
- trainer = Trainer(model, cfg.hyper.train, device)
24
- trainer.train(dataloader, 10)
25
-
26
-
27
- if __name__ == "__main__":
28
- custom_logger()
29
- main()