TinyExp 0.0.1__py3-none-any.whl

tinyexp/__init__.py

@@ -0,0 +1,149 @@
+__author__ = "LI Zeming"
+__email__ = "zane.li@connect.ust.hk"
+__license__ = "MIT"
+
+import os
+from dataclasses import dataclass, field
+from typing import Optional
+
+from hydra.conf import HydraConf, RunDir
+from hydra.core.config_store import ConfigStore
+from omegaconf import DictConfig
+from omegaconf.listconfig import ListConfig
+
+from .exceptions import UnknownConfigurationKeyError
+from .utils.log_utils import tiny_logger_setup
+from .utils.ray_utils import simple_launch_exp
+
+__all__ = ["ConfigStore", "RedisCfgMixin", "TinyExp", "simple_launch_exp"]
+
+
+@dataclass
+class _HydraConfig(HydraConf):
+    """
+    To avoid hydra output the config in unexpected directory.
+    """
+
+    output_subdir: Optional[str] = None
+    run: RunDir = field(default_factory=lambda: RunDir("./output"))
+
+
+@dataclass
+class TinyExp:
+    """
+    Simple experiment configuration class, which use hydra to manage and override configurations.
+    The core idea is to provide a unified interface for experiment configurations, which can be instantiated
+    and used in various contexts, such as Ray or TorchRun.
+    """
+
+    hydra: _HydraConfig = field(default_factory=_HydraConfig)
+
+    # ---------------- luancher configuration ---------------- #
+    num_worker: int = -1  # Number of workers, -1 means to be determined by the user
+    num_gpus_per_worker: float = 1.0  # Number of GPUs per worker, should be a float value between 0 and 1.
+
+    output_root: str = "./output"
+    overrided_cfg: dict = field(default_factory=dict)
+
+    def __repr__(self):
+        # Customize the representation of the Exp object for cleaner Ray logs.
+        return f"Exp(rank={os.getenv('RANK', 'N/A')})"
+
+    @dataclass
+    class WandbCfg:
+        enable_wandb: bool = False
+
+        def build_wandb(self, accelerator=None, **kwargs):
+            if self.enable_wandb:
+                import wandb
+
+                if accelerator is None or accelerator.rank == 0:
+                    wandb.init(**kwargs)
+                return wandb
+
+    wandb_cfg: WandbCfg = field(default_factory=WandbCfg)
+
+    @dataclass
+    class LoggerCfg:
+        def build_logger(self, save_dir: str, distributed_rank: int = 0, filename: str = "log.txt", mode: str = "a"):
+            logger = tiny_logger_setup(save_dir, distributed_rank, filename, mode)
+            logger.info(f"==> log file: {os.path.join(save_dir, filename)}")
+            return logger
+
+    logger_cfg: LoggerCfg = field(default_factory=LoggerCfg)
+
+    def set_cfg(self, cfg_hydra, cfg_object=None):
+        if cfg_object is None:
+            cfg_object = self
+        for key, value in cfg_hydra.items():
+            if hasattr(cfg_object, key):
+                if isinstance(value, (DictConfig, dict)):
+                    # If the value is a dictionary, recursively set attributes
+                    sub_object = getattr(cfg_object, key)
+                    self.set_cfg(value, sub_object)
+                else:
+                    # Otherwise, set the attribute directly
+                    ori_value = getattr(cfg_object, key, None)
+                    if ori_value != value:
+                        if os.getenv("RANK", 0) == 0 or os.getenv("RANK", 0) == "0":
+                            print(f"{key}: {value} <-- {ori_value}(original)")
+                            # print(f"Override {key} from {ori_value} to {value} in {cfg_object.__class__.__name__}")
+                        setattr(cfg_object, key, value)
+                        self.overrided_cfg[key] = value
+            else:
+                raise UnknownConfigurationKeyError(key)
+        return cfg_object
+
+
+@dataclass
+class RedisCfgMixin:
+    @dataclass
+    class RedisCacheCfg:
+        redis_cache_enabled: bool = True
+        redis_cache_shard_ports: ListConfig = field(
+            default_factory=lambda: ListConfig(
+                [
+                    7000,
+                    7001,
+                    7002,
+                    7003,
+                    7004,
+                ]
+            )
+        )  # List of Redis cache shard used ports
+        redis_cache_max_memory: int = 160  # Maximum memory is 160GB, according to the ImageNet dataset size
+        redis_cluster_manager_cpus: int = 10
+
+        def build_redis_cache(self):
+            if self.redis_cache_enabled:
+                from tinyexp.utils.redis_utils import RedisClusterManager
+
+                redis_cluster_manager = RedisClusterManager(
+                    ports=self.redis_cache_shard_ports,
+                    max_memory_per_port=self.redis_cache_max_memory // len(self.redis_cache_shard_ports),
+                )
+                return redis_cluster_manager.start_redis_cluster()
+            return True
+
+    redis_cache_cfg: RedisCacheCfg = field(default_factory=RedisCacheCfg)
+
+    def proxy_build_redis_cache(self):
+        """
+        Hard-coded method to build Redis cache since ray actor need
+        """
+        return self.redis_cache_cfg.build_redis_cache()
+
+
+def store_and_run_exp(exp_class: type[TinyExp]) -> None:
+    """
+    Extract the config from the exp_class and store it in the ConfigStore(hydra config store).
+    Then launch the experiment with the config.
+
+    Args:
+        exp_class: The class of the experiment to run.
+
+    Returns:
+        None: This function does not return anything.
+    """
+    ConfigStore.instance().store(name="cfg", node=exp_class)
+    simple_launch_exp()

tinyexp/dataset/__init__.py

@@ -0,0 +1,7 @@
+# flake8: noqa: F401, F403
+
+from .fake_dataloader import *
+from .sampler import *
+
+_EXCLUDE = {}
+__all__ = [k for k in globals() if k not in _EXCLUDE and not k.startswith("_")]

tinyexp/dataset/fake_dataloader.py

@@ -0,0 +1,25 @@
+import copy
+
+__all__ = ["HoldOnesampleDataLoader"]
+
+
+class HoldOnesampleDataLoader:
+    """
+    A fake dataloader that holds one sample from the original dataloader.
+    This is useful for testing and profiling purposes, where we want to test the model with a single in-memory sample.
+    """
+
+    def __init__(self, dataloader):
+        self.dataloader = dataloader
+        for data in self.dataloader:
+            self.sample = data
+            break
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        return copy.deepcopy(self.sample)
+
+    def __len__(self):
+        return len(self.dataloader)

tinyexp/dataset/sampler.py

@@ -0,0 +1,63 @@
+import itertools
+from typing import Optional
+
+import torch
+from torch.utils.data.sampler import Sampler
+
+__all__ = ["InfiniteSampler"]
+
+
+class InfiniteSampler(Sampler):
+    """
+    In training, we only care about the "infinite stream" of training data.
+    So this sampler produces an infinite stream of indices and
+    all workers cooperate to correctly shuffle the indices and sample different indices.
+    The samplers in each worker effectively produces `indices[worker_id::num_workers]`
+    where `indices` is an infinite stream of indices consisting of
+    `shuffle(range(size)) + shuffle(range(size)) + ...` (if shuffle is True)
+    or `range(size) + range(size) + ...` (if shuffle is False)
+    """
+
+    def __init__(self, size: int, shuffle: bool = True, seed: Optional[int] = 0, drop_last=False, accelerator=None):
+        """
+        Args:
+            size (int): the total number of data of the underlying dataset to sample from
+            shuffle (bool): whether to shuffle the indices or not
+            seed (int): the initial seed of the shuffle. Must be the same
+                across all workers. If None, will use a random seed shared
+                among workers (require synchronization among all workers).
+            drop_last (bool): whether to drop the last incomplete batch
+        """
+        self._size = size
+        self._shuffle = shuffle
+        self._seed = int(seed)
+        self.drop_last = drop_last
+        if accelerator is not None:
+            self._rank = accelerator.rank
+            self._world_size = accelerator.world_size
+        else:
+            self._rank = 0
+            self._world_size = 1
+
+    def set_epoch(self, epoch):
+        pass
+
+    def __iter__(self):
+        start = self._rank
+        yield from itertools.islice(self._infinite_indices(), start, None, self._world_size)
+
+    def _infinite_indices(self):
+        g = torch.Generator()
+        g.manual_seed(self._seed)
+        while True:
+            if self._shuffle:
+                yield from torch.randperm(self._size, generator=g).tolist()
+            else:
+                # yield from torch.arange(self._size)
+                yield from list(range(self._size))
+
+    def __len__(self):
+        if self.drop_last:
+            return self._size // self._world_size
+        else:
+            return (self._size + self._world_size - 1) // self._world_size

tinyexp/examples/mnist_exp.py

@@ -0,0 +1,249 @@
+import datetime
+import os
+from dataclasses import dataclass, field
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+import wandb
+from omegaconf import OmegaConf
+from torch.optim.lr_scheduler import StepLR
+from torchvision import datasets, transforms
+
+from tinyexp import TinyExp, store_and_run_exp
+from tinyexp.exceptions import UnknownAcceleratorTypeError
+
+
+class Net(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.conv1 = nn.Conv2d(1, 32, 3, 1)
+        self.conv2 = nn.Conv2d(32, 64, 3, 1)
+        self.dropout1 = nn.Dropout(0.25)
+        self.dropout2 = nn.Dropout(0.5)
+        self.fc1 = nn.Linear(9216, 128)
+        self.fc2 = nn.Linear(128, 10)
+        self.loss = F.nll_loss
+
+    def forward(self, x, target=None, onnx_exporting=False) -> torch.Tensor:
+        x = self.conv1(x)
+        x = F.relu(x)
+        x = self.conv2(x)
+        x = F.relu(x)
+        x = F.max_pool2d(x, 2)
+        x = self.dropout1(x)
+        x = torch.flatten(x, 1)
+        x = self.fc1(x)
+        x = F.relu(x)
+        x = self.dropout2(x)
+        x = self.fc2(x)
+        if onnx_exporting:
+            return x
+        output = F.log_softmax(x, dim=1)
+
+        if self.training and target is not None:
+            return self.loss(output, target)
+        else:
+            return output
+
+
+@dataclass(repr=False)
+class Exp(TinyExp):
+    exp_class: str = f"{__name__}.Exp"
+    num_worker: int = 2
+    num_gpus_per_worker: float = 0.0
+    mode: str = "train"
+
+    @dataclass
+    class AcceleratorCfg:
+        accelerator: str = "cpu"
+
+        def build_accelerator(self):
+            from tinyexp.tiny_engine.accelerator import CPUAccelerator, DDPAccelerator
+
+            if self.accelerator == "cpu":
+                accelerator = CPUAccelerator()
+            elif self.accelerator == "ddp":
+                accelerator = DDPAccelerator()
+            else:
+                raise UnknownAcceleratorTypeError(self.accelerator)
+            return accelerator
+
+    accelerator_cfg: AcceleratorCfg = field(default_factory=AcceleratorCfg)
+
+    @dataclass
+    class DataloaderCfg:
+        data_root: str = "./data/"
+        train_batch_size_per_device: int = 32
+        train_data_worker_per_gpu: int = 2
+        val_data_worker_per_gpu: int = 1
+
+        def build_train_dataloader(self, accelerator):
+            transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
+            ds_train = datasets.MNIST(self.data_root, train=True, download=True, transform=transform)
+            sampler = torch.utils.data.DistributedSampler(
+                ds_train, num_replicas=accelerator.world_size, rank=accelerator.rank
+            )
+            dl_train = torch.utils.data.DataLoader(
+                ds_train,
+                batch_size=self.train_batch_size_per_device,
+                shuffle=False,
+                num_workers=self.train_data_worker_per_gpu,
+                drop_last=True,
+                sampler=sampler,
+            )
+            return dl_train
+
+        def build_val_dataloader(self, accelerator):
+            transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
+            ds_val = datasets.MNIST(self.data_root, train=False, download=True, transform=transform)
+            sampler = torch.utils.data.DistributedSampler(
+                ds_val, num_replicas=accelerator.world_size, rank=accelerator.rank
+            )
+            dl_val = torch.utils.data.DataLoader(
+                ds_val,
+                batch_size=self.train_batch_size_per_device,
+                shuffle=False,
+                num_workers=self.val_data_worker_per_gpu,
+                drop_last=True,
+                sampler=sampler,
+            )
+            return dl_val
+
+    dataloader_cfg: DataloaderCfg = field(default_factory=DataloaderCfg)
+
+    @dataclass
+    class OptimizerCfg:
+        lr_per_img: float = 1.0 / 64.0  # single image learning rate
+
+        def build_optimizer(self, module, dataloader, accelerator):
+            optimizer = optim.Adadelta(
+                module.parameters(),
+                lr=self.lr_per_img * dataloader.batch_size * accelerator.world_size,
+            )
+            return optimizer
+
+    optimizer_cfg: OptimizerCfg = field(default_factory=OptimizerCfg)
+
+    @dataclass
+    class ModuleCfg:
+        def build_module(self):
+            return Net()
+
+    module_cfg: ModuleCfg = field(default_factory=ModuleCfg)
+
+    @dataclass
+    class LrSchedulerCfg:
+        def build_lr_scheduler(self, optimizer):
+            return StepLR(optimizer, step_size=1, gamma=0.7)
+
+    lr_scheduler_cfg: LrSchedulerCfg = field(default_factory=LrSchedulerCfg)
+
+    # ------------------------------ bellowing is the execution part --------------------- #
+    def run(self) -> None:
+        accelerator = self.accelerator_cfg.build_accelerator()
+        logger = self.logger_cfg.build_logger(
+            save_dir=os.path.join(self.output_root, self.__class__.__name__),
+            distributed_rank=accelerator.rank,
+        )
+        cfg_dict = OmegaConf.to_container(OmegaConf.structured(self), resolve=True)
+        del cfg_dict["hydra"]
+        logger.info(f"-------- Configurations --------\n{OmegaConf.to_yaml(cfg_dict)}")
+
+        self._train(accelerator=accelerator, logger=logger, cfg_dict=cfg_dict)
+
+    def _evaluate(self, accelerator, logger, module_or_module_path, val_dataloader=None) -> None:
+        if isinstance(module_or_module_path, str):
+            module = Net()
+            module.load_state_dict(torch.load(module_or_module_path))
+            module = accelerator.prepare(module)
+        else:
+            module = module_or_module_path
+
+        if val_dataloader is None:
+            val_dataloader = self.dataloader_cfg.build_val_dataloader(accelerator)
+
+        module.eval()
+        accurate = torch.tensor(0.0, device=accelerator.device)
+
+        for batch in val_dataloader:
+            features, labels = (item.to(accelerator.device) for item in batch)
+            with torch.no_grad():
+                preds = module(features)
+            predictions = preds.argmax(dim=-1)
+            accurate_preds = predictions == labels
+            accurate_preds_sum = accelerator.reduce_sum(accurate_preds.sum())
+            accurate += accurate_preds_sum
+        eval_metric = accurate.item() / len(val_dataloader.dataset)
+
+        accelerator.wait_for_everyone()
+        nowtime = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+        logger.info(f"{nowtime} --> eval_metric= {100 * eval_metric:.2f}%")
+
+        if self.wandb_cfg.enable_wandb and accelerator.is_main_process:
+            wandb.log({"val_metric": eval_metric})
+
+    def _train(self, accelerator, logger, cfg_dict) -> None:
+        train_dataloader = self.dataloader_cfg.build_train_dataloader(accelerator)
+        val_dataloader = self.dataloader_cfg.build_val_dataloader(accelerator)
+        ori_module = self.module_cfg.build_module()
+        ori_optimizer = self.optimizer_cfg.build_optimizer(ori_module, train_dataloader, accelerator)
+        lr_scheduler = self.lr_scheduler_cfg.build_lr_scheduler(ori_optimizer)
+
+        module, optimizer = accelerator.prepare(ori_module, ori_optimizer)
+        accelerator.print(f"device {accelerator.device!s} is used!")
+
+        train_iter = iter(train_dataloader)
+        if self.wandb_cfg.enable_wandb and accelerator.rank == 0:
+            self.wandb_cfg.build_wandb(
+                accelerator=accelerator,
+                project="Baselines",
+                config=cfg_dict,
+            )
+
+        for epoch in range(3):
+            module.train()
+
+            for step in range(len(train_dataloader)):
+                try:
+                    batch = next(train_iter)
+                except StopIteration:
+                    train_iter = iter(train_dataloader)
+                    batch = next(train_iter)
+
+                features, labels = (item.to(accelerator.device) for item in batch)
+                preds = module(features)
+                loss = nn.CrossEntropyLoss()(preds, labels)
+
+                optimizer.zero_grad()
+                accelerator.backward(loss)
+                optimizer.step()
+                if (step + 1) % 20 == 0:
+                    logger.info(f"epoch {epoch} loss: {loss.item(): .4f} lr: {optimizer.param_groups[0]['lr']: .4f}")
+                    if self.wandb_cfg.enable_wandb and accelerator.rank == 0:
+                        wandb.log(
+                            {
+                                "epoch": epoch,
+                                "loss": loss.item(),
+                                "lr": optimizer.param_groups[0]["lr"],
+                            }
+                        )
+            self._evaluate(
+                accelerator=accelerator, logger=logger, module_or_module_path=module, val_dataloader=val_dataloader
+            )
+
+            lr_scheduler.step()
+
+
+# import hydra
+# @hydra.main(version_base=None, config_name="cfg")
+# def simple_launch_exp(cfg: DictConfig) -> None:
+#     from omegaconf import DictConfig, OmegaConf
+#     print(OmegaConf.to_yaml(cfg))
+#     from IPython import embed; embed()  # for debugging
+#     exp_class = hydra.utils.get_class(cfg.exp_class)
+#     exp_class().set_cfg(cfg).run()
+
+if __name__ == "__main__":
+    store_and_run_exp(Exp)