🔥 [Remove] train.py which outside of ./yolo
Browse files
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()
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|