moskomule / chika

chika is a simple and easy config tool for hierarchical configurations.
MIT License
20 stars 3 forks source link

chika

chika is a simple and easy config tool for hierarchical configurations.

Requirements

Installation

pip install -U chika

or

pip install -U git+https://github.com/moskomule/chika

Usage

Write typed configurations using chika.config, which depends on dataclass.

# main.py
import enum
import chika

@chika.config
class ModelConfig:
    name: str = chika.choices('resnet', 'densenet')
    depth: int = 10

@chika.config
class DataConfig:
    # values that needs to be specified
    name: str = chika.required(help="name of dataset")

class Optims(str, enum.Enum):
    # currently, only StrEnum is supported
    sgd = "sgd"
    adam = "adam"

@chika.config
class OptimConfig:
    name: Optims = Optims.sgd
    steps: list[int] = chika.sequence(100, 150)

@chika.config
class BaseConfig:
    model: ModelConfig
    data: DataConfig
    optim: OptimConfig

    seed: int = chika.with_help(1, help="random seed")
    use_amp: bool = False
    gpu: int = chika.choices(*range(torch.cuda.device_count()), help="id of gpu")

Then, wrap the main function with chika.main(BaseConfig).

@chika.main(BaseConfig)
def main(cfg: BaseConfig):
    set_seed(cfg.seed)
    model = ModelRegistry(cfg.model)
    ...

Now, main.py can be executed as...

python main.py --use_amp
# cfg.use_amp == True

python main.py --model config/densenet.yaml
# cfg.model.name == densenet
# cfg.model.depth == 12

python main.py --model config/densenet.yaml --model.depth 24
# cfg.model.name == densenet
# cfg.model.depth == 24

python main.py --optim.decay_steps 120 150
# config.optim.decay_steps == [120, 150]

python main.py -h
#usage: test.py [-h] [--model MODEL] [--model.name {resnet,densenet}] [--model.depth MODEL.DEPTH] [--data DATA] --data.name DATA.NAME [--optim OPTIM] [--optim.steps OPTIM.STEPS [OPTIM.STEPS ...]]
#               [--seed SEED] [--use_amp] [--gpu {1,2,3}]
#
#optional arguments:
#  -h, --help            show this help message and exit
#  --model MODEL         load {yaml,yml,json} file for model if necessary
#  --model.name {resnet,densenet}
#                        (default: 'resnet')
#  --model.depth MODEL.DEPTH
#                        (default: 10)
#  --data DATA           load {yaml,yml,json} file for data if necessary
#  --data.name DATA.NAME
#                        name of dataset (required) (default: None)
#  --optim OPTIM         load {yaml,yml,json} file for optim if necessary
#  --optim.steps OPTIM.STEPS [OPTIM.STEPS ...]
#                        (default: [100, 150])
#  --seed SEED           random seed (default: 1)
#  --use_amp             (default: False)
#  --gpu {1,2,3}         id of gpu (default: 1)

Child configs can be updated via YAML or JSON files.

An example of YAML file (e.g. config/densenet.yaml)

name: densenet
depth: 12 

An example of JSON file (e.g. config/densenet.json)

{
  "name": "densenet",
  "depth": 12
}

For chika.Config, the following functions are prepared:

def with_help(default, help): ...

# add help message
def choices(*values, help): ...

# add candidates that should be selected
def sequence(*values, size, help): ...

# add a list. size can be specified
def required(*, help): ...

# add a required value
def bounded(default, _from, _to, *help): ...
# add boundaries

Working Directory

change_job_dir=True creates a unique directory for each run.

@chika.main(BaseConfig, change_job_dir=True)
def main(cfg):
    print(Path(".").resolve())
    # /home/user/outputs/0901-122412-558f5a
    print(Path(".") / "weights.pt")
    # /home/user/outputs/0901-122412-558f5a/weights.pt
    print(chika.original_path)
    # /home/user
    print(chika.resolve_original_path("weights.pt"))
    # /home/user/weights.pt

Other APIs

from chika import ChikaConfig

cfg = ChikaConfig.from_dict(...)

cfg.to_dict()
# {"model": {"name": "resnet", "zero_init": True, ...}, ...}

Known issues and ToDos