yukw777

[tool.poetry]
name = "test"
version = "0.1.0"
description = ""
authors = ["Peter Yu <2057325+yukw777@users.noreply.github.com>"]
readme = "README.md"

[tool.poetry.dependencies]
python = "~3.9"

yukw777

@pytest.mark.parametrize("network_size", ["small", "big", "huge"])
def test_train_network_size(network_size):
    with initialize(config_path="../leela_zero_pytorch/conf"):
        cfg = compose(
            config_name="config",
            overrides=[
                f"+network={network_size}",
                "data.train_data_dir=tests/test-data",
                "data.train_dataloader_conf.batch_size=2",
                "data.val_data_dir=tests/test-data",

yukw777

class NetworkLightningModule(..., pl.LightningModule):
    def __init__(self, ...):
        super().__init__(...)
        self.save_hyperparameters()

        # metrics
        self.train_accuracy = pl.metrics.Accuracy()
        self.val_accuracy = pl.metrics.Accuracy()
        self.test_accuracy = pl.metrics.Accuracy()

yukw777

def training_step(self, batch: DataPoint, batch_idx: int) -> torch.Tensor:
    planes, target_move, target_val = batch
    pred_move, pred_val = self(planes)
    mse_loss, cross_entropy_loss, loss = self.loss(
        pred_move, pred_val, target_move, target_val
    )
    self.log("train_loss", loss, prog_bar=True)
    self.log_dict(
        {
            "train_mse_loss": mse_loss,

yukw777

import logging
import hydra

from omegaconf import DictConfig, OmegaConf
from pytorch_lightning import Trainer
from hydra.utils import instantiate


logger = logging.getLogger(__name__)

yukw777

@pytest.mark.parametrize("network_size", ["small", "big", "huge"])
def test_train_network_size(monkeypatch, tmp_path, capsys, network_size):
    with initialize(config_path="../leela_zero_pytorch/conf"):
        cfg = compose(
            config_name="config",
            overrides=[
                f"+network={network_size}",
                "data.train_data_dir=tests/test-data",
                "data.train_dataloader_conf.batch_size=2",
                "data.val_data_dir=tests/test-data",

yukw777

@hydra.main(config_path="conf", config_name="config")
def main(cfg: DictConfig) -> Trainer:
    logger.info(f"Training with the following config:\n{OmegaConf.to_yaml(cfg)}")

    network = instantiate(cfg.network, cfg.train)

yukw777

# @package _group_
_target_: leela_zero_pytorch.network.NetworkLightningModule
network_conf:
  residual_channels: 128
  residual_layers: 6

yukw777

# @package _group_
_target_: leela_zero_pytorch.network.NetworkLightningModule
network_conf:
  residual_channels: 32
  residual_layers: 8

yukw777

class DataModule(pl.LightningDataModule):
    def __init__(
        self,
        train_data_dir: str,
        val_data_dir: str,
        test_data_dir: str,
        train_dataloader_conf: Optional[DictConfig] = None,
        val_dataloader_conf: Optional[DictConfig] = None,
        test_dataloader_conf: Optional[DictConfig] = None,
    ):