diff --git a/.gitignore b/.gitignore index 6e5c068..0b4eae7 100644 --- a/.gitignore +++ b/.gitignore @@ -8,6 +8,7 @@ __pycache__/ # Folder **/PokemonData /runs +outputs # Weights and optimizer weights/model.pt diff --git a/README.md b/README.md index 5000718..451cb0a 100644 --- a/README.md +++ b/README.md @@ -88,7 +88,12 @@ Example script to train and evaluate model. python main.py ``` +**Hydra - hyperparameter management tool.** + +If you want to change hyper-parameters such as: `epoch_count`, `lr`, `batch_size`, `momentum`, +then you can change them in the file `conf/config.yaml`. + ## (A) Acknowledgments This repository borrows partially from [Isadrtdinov](https://github.com/isadrtdinov/intro-to-dl-hse/blob/2022-2023/seminars/201/seminar_04.ipynb), and [FUlyankin](https://github.com/FUlyankin/deep_learning_pytorch/tree/main/week08_fine_tuning) repositories. -Repository design taken from [v-goncharenko](https://github.com/v-goncharenko/data-science-template) and [PeterWang512](https://github.com/PeterWang512/CNNDetection). +Repository design taken from [v-goncharenko](https://github.com/v-goncharenko/data-science-template), [PeterWang512](https://github.com/PeterWang512/CNNDetection) and [ArjanCodes](https://github.com/ArjanCodes/2021-config). diff --git a/conf/config.yaml b/conf/config.yaml new file mode 100644 index 0000000..7c0c0d7 --- /dev/null +++ b/conf/config.yaml @@ -0,0 +1,9 @@ +defaults: + - _self_ +paths: + data: ${hydra:runtime.cwd}/data/PokemonData +params: + epoch_count: 10 + lr: 1e-2 + batch_size: 32 + momentum: 0.9 diff --git a/config.py b/config.py new file mode 100644 index 0000000..9ff8812 --- /dev/null +++ b/config.py @@ -0,0 +1,20 @@ +from dataclasses import dataclass + + +@dataclass +class Paths: + data: str + + +@dataclass +class Params: + epoch_count: int + lr: float + batch_size: int + momentum: float + + +@dataclass +class ConvNetConfig: + paths: Paths + params: Params diff --git a/ds/dataset.py b/ds/dataset.py index bd65fad..c316bcc 100644 --- a/ds/dataset.py +++ b/ds/dataset.py @@ -10,8 +10,9 @@ from tqdm import tqdm -def remove_bed_images(data_dir: Path): - bad_images = glob.glob(f"{data_dir}/*/*.svg") +def remove_bed_images(root_path: str): + data_path = Path(f"{root_path}") + bad_images = glob.glob(f"{data_path}/*/*.svg") for bad_image in bad_images: os.remove(bad_image) @@ -109,21 +110,21 @@ def __getitem__(self, item): def create_dataloader( - root: Path, + root_path: str, batch_size: int, load_to_ram: bool = False, pin_memory: bool = True, num_workers: int = 2, ) -> tuple[DataLoader[Any], DataLoader[Any]]: train_dataset = PokemonDataset( - root=root, + root=Path(f"{root_path}"), train=True, load_to_ram=load_to_ram, transform=prepare_train_data(), ) test_dataset = PokemonDataset( - root=root, + root=Path(f"{root_path}"), train=False, load_to_ram=load_to_ram, transform=prepare_test_data(), diff --git a/main.py b/main.py index f4cfe53..4f10a67 100644 --- a/main.py +++ b/main.py @@ -1,40 +1,44 @@ -import pathlib - +import hydra import torch +from hydra.core.config_store import ConfigStore +from omegaconf import OmegaConf +from config import ConvNetConfig from ds.dataset import create_dataloader, remove_bed_images from ds.models import ConvNet from ds.runner import train -# Hyper parameters -EPOCH_COUNT = 10 -LR = 1e-2 -MOMENTUM = 0.9 -BATCH_SIZE = 32 - -# Data configuration -DATA_DIR = pathlib.Path("data/PokemonData") - # Device device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') -def main(): +cs = ConfigStore.instance() +cs.store(name="ConfNet_config", node=ConvNetConfig) + + +@hydra.main(config_path="conf", config_name="config") +def main(cfg: ConvNetConfig): + print(OmegaConf.to_yaml(cfg)) + # Model and Optimizer model_name = "ConvNet" model = ConvNet().to(device) - optimizer = torch.optim.SGD(model.parameters(), LR, MOMENTUM) + optimizer = torch.optim.SGD( + model.parameters(), lr=cfg.params.lr, momentum=cfg.params.momentum + ) criterion = torch.nn.CrossEntropyLoss() - scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, EPOCH_COUNT) + scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( + optimizer=optimizer, T_max=cfg.params.epoch_count + ) # Remove bad images - remove_bed_images(DATA_DIR) + remove_bed_images(cfg.paths.data) # Create the data loaders train_loader, test_loader = create_dataloader( - root=DATA_DIR, - batch_size=BATCH_SIZE, + root_path=cfg.paths.data, + batch_size=cfg.params.batch_size, load_to_ram=False, pin_memory=True, num_workers=2, @@ -48,7 +52,7 @@ def main(): criterion, train_loader, test_loader, - num_epochs=EPOCH_COUNT, + num_epochs=cfg.params.epoch_count, device=device, title=model_name, ) diff --git a/poetry.lock b/poetry.lock index 2d7e443..30b5991 100644 --- a/poetry.lock +++ b/poetry.lock @@ -4643,4 +4643,4 @@ test = ["zope.testing"] [metadata] lock-version = "2.0" python-versions = "^3.11" -content-hash = "e76bca6785924e9893cfee608aaf0229c409c9d728bbd77002dfcee51f2c2cc4" +content-hash = "8e31ae4b4922d37853b0c1d85cac45f30601afccf62f5d25e7f7cfa21a2282ae" diff --git a/pyproject.toml b/pyproject.toml index e9170b6..2d01277 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -21,6 +21,7 @@ isort = "^5.13.0" flake8 = "^6.1.0" pre-commit = "^3.6.0" dvc = {extras = ["gdrive"], version = "^3.33.4"} +hydra-core = "^1.3.2" [tool.black] line-length = 90 diff --git a/weights/download_weights.sh b/weights/download_weights.sh index 89c8d0f..24b51a0 100644 --- a/weights/download_weights.sh +++ b/weights/download_weights.sh @@ -1,2 +1,2 @@ -gdown --id 1R7NCduqaDK_D2R4mgUSvgZubcX5DjSYr -O model.pt -gdown --id 1Ds3Hch72xXUyEV0_aiV8TwNJp8Pg3J9H -O optimizer.pt +gdown --id 1ROwjrRiM_EfqegJEl2RxW3pM9ieK80Cv -O model.pt +gdown --id 11FzTrBeYoj3MQ8bmnEZ6BDP6JtXjFylz -O optimizer.pt