chugr 0.2.0.dev0__py3-none-any.whl

chugr/common/config.py

@@ -0,0 +1,335 @@
+from dataclasses import dataclass, field, fields, replace
+from typing import Any, Dict, List, Optional, Tuple, Union, Sequence, Callable
+
+from .serialize import Serializable
+
+from .types import ImageFeatureInfo, ShardSpec, SourceSpec
+from .urls import expand_urls
+
+
+def image_mode_to_chs(fmt: str):
+    if fmt is None:
+        return None
+    assert fmt in ('L', 'RGB')   # could support more...
+    return 1 if fmt == 'L' else 3
+
+
+@dataclass
+class ImageInputCfg(Serializable):
+    size: Optional[Tuple[int, int]] = (512, 512)
+    mode: Optional[str] = 'L'
+    mean: Optional[Union[float, Tuple[float, ...]]] = 0.5
+    std: Optional[Union[float, Tuple[float, ...]]] = 0.5
+    interpolation: Optional[str] = 'bicubic'
+    fill_color: Optional[Union[int, Tuple[int, ...]]] = 255
+    crop_margin: Optional[bool] = False
+    align_long_axis: Optional[bool] = False
+    transform_type: Optional[str] = 'image_basic'
+    resize_mode: Optional[str] = 'shortest'
+    patch_size: Optional[int] = None  # patch size for naflex (variable aspect / resolution) transforms
+    max_seq_len: Optional[int] = None  # max patch sequence length for naflex transforms
+
+    @property
+    def image_chs(self):
+        return image_mode_to_chs(self.mode)
+
+    def __post_init__(self):
+        image_chs = self.image_chs
+        if image_chs is not None:
+            # ensue mean/std attributes match # image_chs
+            for attr_name in ('mean', 'std'):
+                attr = getattr(self, attr_name)
+                if attr is not None and not isinstance(attr, Sequence):
+                    attr = (attr,)
+                if image_chs == 1 and len(attr) > image_chs:
+                    attr = (sum(attr) / len(attr),)
+                if image_chs > 1 and len(attr) == 1:
+                    attr = attr * image_chs
+                assert len(attr) == image_chs
+                setattr(self, attr_name, attr)
+
+            # ensure fill color matches image_chs
+            if self.fill_color is not None:
+                if not isinstance(self.fill_color, Sequence):
+                    self.fill_color = (self.fill_color,)
+                if image_chs == 1 and len(self.fill_color) > image_chs:
+                    self.fill_color = (int(sum(self.fill_color) / len(self.fill_color)),)
+                if image_chs > 1 and len(self.fill_color) == 1:
+                    self.fill_color  = self.fill_color * image_chs
+
+    @classmethod
+    def empty(cls):
+        return cls(**{f.name: None for f in fields(cls)})
+
+    def set_default(self, right, inplace=False):
+        # set left fields from right fields if left fields are not-initialized (None)
+        changes = {
+            f.name: v for f in fields(self)
+            if (v := getattr(right, f.name)) is not None and getattr(self, f.name) is None
+        }
+        if inplace:
+            for k, v in changes.items():
+                setattr(self, k, v)
+            return self
+        else:
+            return replace(self, **changes)
+
+    def merge(self, right, inplace=False):
+        # merge from right to left for right fields that are not None
+        changes = {f.name: v for f in fields(self) if (v := getattr(right, f.name)) is not None}
+        if inplace:
+            for k, v in changes.items():
+                setattr(self, k, v)
+            return self
+        else:
+            return replace(self, **changes)
+
+
+@dataclass
+class ImageAugCfg(Serializable):
+    """
+    A simple flat config struct for overriding common augmentation defaults.
+
+    Each image transform type supports different augmentations and have their own defaults,
+    this struct is intended to override defaults for common values, not necessarily to
+    cover all cases and define all augmentation possibilities across all schemes.
+    """
+
+    # resize scale bounds (1.0 = middle point = same scale)
+    scale: Optional[Tuple[float, float]] = None
+
+    # resize aspect ratio bounds (1.0 = 1:1)
+    ratio: Optional[Tuple[float, float]] = None
+
+    # color jitter, per item probs
+    color_jitter: Optional[Union[float, Tuple[float, float, float], Tuple[float, float, float, float]]] = None
+
+    # for simclr, control prob for applying any of the jitter probs above
+    color_jitter_prob: Optional[float] = None
+
+    # for preprocess w/ grayscale (simclr), control prob of converting to graysacle
+    grayscale_prob: Optional[float] = None
+
+    gaussian_blur_prob: Optional[float] = None
+    gaussian_blur_kernel_size: Optional[int] = None
+
+    # probability of applying random-erasing (timm style aug)
+    re_prob: Optional[float] = None
+
+    # number of random-erasing blocks (timm style aug)
+    re_count: Optional[int] = None
+
+    @classmethod
+    def clip(cls, **kwargs):
+        aug = cls(
+            scale=(0.9, 1.0),
+            ratio=(0.75, 1. / 0.75),
+        )
+        aug = replace(aug, **kwargs)
+        return aug
+
+    @classmethod
+    def imagenet(cls, **kwargs):
+        aug = cls(
+            scale=(0.08, 1.0),
+            ratio=(0.75, 1. / 0.75),
+            color_jitter=(0.4, 0.4, 0.4),
+        )
+        aug = replace(aug, **kwargs)
+        return aug
+
+    @classmethod
+    def simclr(cls, **kwargs):
+        aug = cls(
+            scale=(0.08, 1.0),
+            ratio=(0.75, 1. / 0.75),
+            color_jitter=(0.4, 0.4, 0.4, 0.1),
+            color_jitter_prob=0.8,
+            grayscale_prob=0.2,
+            gaussian_blur_prob=0.5,
+            #gaussian_blur_kernel_size=23,
+        )
+        aug = replace(aug, **kwargs)
+        return aug
+
+
+# Vision preprocessing config
+@dataclass
+class PreprocessCfg(Serializable):
+    image_input: ImageInputCfg = field(default_factory=ImageInputCfg)
+    aug_cfg: Optional[ImageAugCfg] = None
+
+
+@dataclass
+class DataArg(Serializable):
+    """ Data source argument in an argument friendly form (multiple sources represented in string)
+    """
+    source: str
+    split: Optional[str] = None
+    sampling_weight: Optional[str] = None
+    template: Optional[str] = None  # template to transform url for use
+    num_samples: Optional[Union[int, str]] = None
+    data_dir: Optional[str] = None
+
+    batch_size: int = 1
+    format: str = "wds"  # e.g. "hfds", "hfids", or "wds"
+
+    resampled: bool = False  # sample shards with replacement
+    multi_interval: bool = True
+    persistent_workers: bool = True
+    num_workers: int = 4
+
+    # NaFlex schedule driven batching (variable resolution buckets w/ token budget batch sizes), enabled when
+    # naflex_seq_lens is set (batch_size is then ignored for training).
+    naflex_seq_lens: Optional[List[int]] = None
+    naflex_seq_len_probs: Optional[List[float]] = None
+    naflex_max_tokens: Optional[int] = None  # tokens (patches + text) per batch per rank
+    naflex_text_tokens: int = 0  # per row text token cost added when sizing batches
+    naflex_batch_divisor: int = 8
+    length_bucketing: bool = False  # pool sort samples by annotation text length
+
+
+def split_sources(
+        source: str,
+        split: Optional[str] = None,
+        sampling_weights: Optional[Union[str, List[float]]] = None,
+        num_samples: Optional[Union[int, str, List[int]]] = None,
+):
+    if '::' in source:
+        source_split = source.split('::')
+    else:
+        source_split = [source]
+    num_sources = len(source_split)
+
+    if sampling_weights is not None:
+        if isinstance(sampling_weights, str):
+            weights_split = sampling_weights.split('::')
+            sampling_weights = [float(w) for w in weights_split]
+        assert len(sampling_weights) == num_sources
+
+    num_samples_per_source = None
+    if num_samples is not None:
+        if isinstance(num_samples, str):
+            num_samples_split = num_samples.split('::')
+            num_samples = [int(s) for s in num_samples_split]
+
+        try:
+            len(num_samples)
+        except Exception:
+            num_samples_per_source = [None] * num_sources
+        else:
+            num_samples_per_source = num_samples
+            num_samples = sum(num_samples_per_source)
+        finally:
+            assert len(num_samples_per_source) == num_sources
+
+    output = []
+    for i, s in enumerate(source_split):
+        output.append(SourceSpec(
+            url=s,
+            split=split,
+            sampling_weight=None if sampling_weights is None else sampling_weights[i],
+            num_samples=None if num_samples_per_source is None else num_samples_per_source[i],
+        ))
+
+    return output, num_samples
+
+
+# FIXME add code to resolve shard information from _info.yaml or .json files (see dataset_info.py)
+
+
+def source_to_shard_spec(
+        source: Union[str, SourceSpec, List[SourceSpec]],
+):
+    if isinstance(source, str):
+        source_list = [SourceSpec(url=source)]
+    elif isinstance(source, SourceSpec):
+        source_list = [source]
+    else:
+        assert isinstance(source[0], SourceSpec)
+        source_list = source
+
+    # process weights first in case some are set and some are not
+    if not all(s.sampling_weight is None for s in source_list):
+        weights = [s.sampling_weight if s.sampling_weight else 1.0 for s in source_list]
+    else:
+        weights = [None] * len(source_list)
+
+    all_urls = []
+    all_weights = []
+    for s, w in zip(source_list, weights):
+        expanded_urls, expanded_weights = expand_urls(s.url, weights=w)
+        all_urls.extend(expanded_urls)
+        if expanded_weights:
+            all_weights.extend(expanded_weights)
+    all_weights = all_weights or None
+    sizes = None  # FIXME resolve sizes
+
+    ss = ShardSpec(urls=all_urls, weights=all_weights, sizes=sizes)
+    return ss
+
+
+@dataclass
+class DataCfg(Serializable):
+    source: Union[str, SourceSpec, List[SourceSpec]]
+    split: Optional[str] = None
+    num_samples: Optional[int] = None  # overrides num_samples across sources if set
+    data_dir: Optional[str] = None
+
+    batch_size: Optional[int] = 1
+    format: str = "wds"  # e.g. "hfds", "hfids", or "wds".
+
+    resampled: bool = False  # sample shards with replacement
+    multi_interval: bool = True
+    persistent_workers: bool = True
+    num_workers: int = 4
+
+    # NaFlex schedule driven batching (variable resolution buckets w/ token budget batch sizes), enabled when
+    # naflex_seq_lens is set (batch_size is then ignored for training).
+    naflex_seq_lens: Optional[List[int]] = None
+    naflex_seq_len_probs: Optional[List[float]] = None
+    naflex_max_tokens: Optional[int] = None  # tokens (patches + text) per batch per rank
+    naflex_text_tokens: int = 0  # per row text token cost added when sizing batches
+    naflex_batch_divisor: int = 8
+    length_bucketing: bool = False  # pool sort samples by annotation text length
+
+    @classmethod
+    def from_arg(cls, data_arg: DataArg):
+        sources, num_samples = split_sources(
+            data_arg.source,
+            split=data_arg.split,
+            sampling_weights=data_arg.sampling_weight,
+            num_samples=data_arg.num_samples,
+        )
+        return cls(
+            source=sources if len(sources) > 1 else sources[0],
+            split=data_arg.split,
+            num_samples=num_samples,
+            data_dir=data_arg.data_dir,
+            batch_size=data_arg.batch_size,
+            format=data_arg.format,
+            resampled=data_arg.resampled,
+            multi_interval=data_arg.multi_interval,
+            persistent_workers=data_arg.persistent_workers,
+            num_workers=data_arg.num_workers,
+            naflex_seq_lens=data_arg.naflex_seq_lens,
+            naflex_seq_len_probs=data_arg.naflex_seq_len_probs,
+            naflex_max_tokens=data_arg.naflex_max_tokens,
+            naflex_text_tokens=data_arg.naflex_text_tokens,
+            naflex_batch_divisor=data_arg.naflex_batch_divisor,
+            length_bucketing=data_arg.length_bucketing,
+        )
+
+    @property
+    def shard_spec(self):
+        return source_to_shard_spec(self.source)
+
+    def __post_init__(self):
+        if self.num_workers == 0:
+            self.persistent_workers = False
+
+@dataclass
+class DistributedCfg:
+    world_size: int = 1
+    local_rank: int = 0
+    global_rank: int = 0

chugr/common/random.py

@@ -0,0 +1,76 @@
+import random
+
+import numpy as np
+
+from .types import SharedCount
+
+
+_RANK_SEED_STRIDE = 1000003
+
+
+def rank_seed(seed: int, rank: int = 0) -> int:
+    """ Seed for per rank randomness (process RNGs, dataloader worker seeds).
+
+    Ranks are spaced by a large stride (vs seed + rank) so runs w/ nearby seeds don't share random streams
+    (seed 42 rank 1 == seed 43 rank 0), rank 0 keeps `seed` so single process runs are unchanged.
+    """
+    return seed + _RANK_SEED_STRIDE * rank
+
+
+def seed_worker(worker_id):
+    import torch
+    worker_seed = torch.initial_seed()
+    random.seed(worker_seed)
+    np.random.seed(worker_seed % 2**32)
+
+
+def get_pytorch_worker_seed(increment=0, initial_seed=None):
+    """get dataloader worker seed from pytorch
+    """
+    from torch.utils.data import get_worker_info
+
+    increment_value = increment.get_value() if isinstance(increment, SharedCount) else increment
+    worker_info = get_worker_info()
+    if worker_info is not None:
+        # favour using the seed already created for pytorch dataloader workers if it exists
+        seed = worker_info.seed
+        num_workers = worker_info.num_workers
+        if increment_value:
+            # space out seed increments so they can't overlap across workers in different iterations
+            seed += increment_value * max(1, num_workers)
+    else:
+        # a fallback when no dataloader workers are present (num_workers=0)
+        import torch
+
+        if initial_seed is None:
+            initial_seed = torch.initial_seed()
+
+        # generate seed from initial via torch.Generator so it matches DL worker seeds
+        seed = torch.empty((), dtype=torch.int64).random_(
+            generator=torch.Generator().manual_seed(initial_seed)).item()
+
+        if increment_value:
+            seed += increment_value
+
+    return seed
+
+
+def shuffled_interval_indices(dataset_len: int, num_indices: int, interval: int = 0, seed: int = 0):
+    """ Shuffled dataset indices for one interval of a map-style dataset (all ranks, split w/ [rank::world_size]).
+
+    Each pass over the dataset is a fresh permutation (seed + pass), intervals consume the permutations
+    sequentially (interval i starts where i - 1 ended), so intervals shorter than the dataset still cycle through
+    every sample before repeating. Identical across ranks / workers for the same args.
+    """
+    import torch
+    assert dataset_len > 0, 'dataset is empty'
+    epoch, offset = divmod(interval * num_indices, dataset_len)
+    indices = []
+    while len(indices) < num_indices:
+        g = torch.Generator()
+        g.manual_seed(seed + epoch)
+        perm = torch.randperm(dataset_len, generator=g).tolist()
+        indices.extend(perm[offset:])
+        offset = 0
+        epoch += 1
+    return indices[:num_indices]

chugr/common/serialize.py

@@ -0,0 +1,150 @@
+""" Minimal dataclass (de)serialization
+
+A small replacement for the subset of `simple_parsing.helpers.Serializable` used by chugr & pixparse.
+Supports nested dataclasses, Optional / Union / Tuple / List / Dict fields, and json or yaml files.
+"""
+import dataclasses
+import json
+import types
+import typing
+from pathlib import Path
+from typing import Any, Dict, Type, TypeVar, Union
+
+T = TypeVar('T')
+
+_NONE_TYPE = type(None)
+
+
+def _is_union(tp) -> bool:
+    origin = typing.get_origin(tp)
+    return origin is Union or origin is types.UnionType
+
+
+def _unwrap_optional(tp):
+    """ Optional[X] -> X, leave other unions untouched """
+    if _is_union(tp):
+        args = [a for a in typing.get_args(tp) if a is not _NONE_TYPE]
+        if len(args) == 1:
+            return args[0]
+    return tp
+
+
+def resolve_field_types(cls) -> Dict[str, Any]:
+    try:
+        return typing.get_type_hints(cls)
+    except Exception:
+        return {f.name: f.type for f in dataclasses.fields(cls)}
+
+
+def to_dict(obj) -> Any:
+    """ Recursively convert dataclasses (and containers of) to plain python types. """
+    if dataclasses.is_dataclass(obj) and not isinstance(obj, type):
+        out = {}
+        for f in dataclasses.fields(obj):
+            if not f.init and not f.metadata.get('serialize', False):
+                continue
+            out[f.name] = to_dict(getattr(obj, f.name))
+        return out
+    if isinstance(obj, (list, tuple)):
+        return [to_dict(v) for v in obj]
+    if isinstance(obj, dict):
+        return {k: to_dict(v) for k, v in obj.items()}
+    if isinstance(obj, Path):
+        return str(obj)
+    return obj
+
+
+def from_value(tp, value):
+    """ Convert a plain python value to the annotated type `tp`, recursing into dataclasses. """
+    if value is None:
+        return None
+    tp = _unwrap_optional(tp)
+
+    if dataclasses.is_dataclass(tp):
+        if isinstance(value, tp):
+            return value
+        assert isinstance(value, dict), f'Expected dict for {tp}, got {type(value)}'
+        return from_dict(tp, value)
+
+    if _is_union(tp):
+        # try dataclass members first when given a dict, otherwise pass through unchanged
+        if isinstance(value, dict):
+            for a in typing.get_args(tp):
+                if dataclasses.is_dataclass(a):
+                    try:
+                        return from_dict(a, value)
+                    except (TypeError, AssertionError):
+                        continue
+        if isinstance(value, list):
+            for a in typing.get_args(tp):
+                if typing.get_origin(a) in (list, tuple):
+                    return from_value(a, value)
+        return value
+
+    origin = typing.get_origin(tp)
+    args = typing.get_args(tp)
+    if origin is tuple:
+        if not isinstance(value, (list, tuple)):
+            value = (value,)
+        if len(args) == 2 and args[1] is Ellipsis:
+            return tuple(from_value(args[0], v) for v in value)
+        if args and len(args) == len(value):
+            return tuple(from_value(a, v) for a, v in zip(args, value))
+        return tuple(value)
+    if origin is list:
+        item_tp = args[0] if args else Any
+        return [from_value(item_tp, v) for v in value]
+    if origin is dict:
+        val_tp = args[1] if len(args) == 2 else Any
+        return {k: from_value(val_tp, v) for k, v in value.items()}
+    if tp is float and isinstance(value, int) and not isinstance(value, bool):
+        return float(value)
+    return value
+
+
+def from_dict(cls: Type[T], data: Dict[str, Any]) -> T:
+    field_types = resolve_field_types(cls)
+    init_fields = {f.name for f in dataclasses.fields(cls) if f.init}
+    unknown = set(data.keys()) - init_fields
+    if unknown:
+        raise TypeError(f'Unknown fields {sorted(unknown)} for {cls.__name__}.')
+    kwargs = {k: from_value(field_types[k], v) for k, v in data.items()}
+    return cls(**kwargs)
+
+
+def _load_file(path: Union[str, Path]) -> Dict[str, Any]:
+    path = Path(path)
+    text = path.read_text()
+    if path.suffix in ('.yaml', '.yml'):
+        import yaml
+        return yaml.safe_load(text) or {}
+    return json.loads(text)
+
+
+class Serializable:
+    """ Mixin for dataclasses w/ dict, json, and yaml (de)serialization. """
+
+    def to_dict(self) -> Dict[str, Any]:
+        return to_dict(self)
+
+    @classmethod
+    def from_dict(cls: Type[T], data: Dict[str, Any]) -> T:
+        return from_dict(cls, data)
+
+    @classmethod
+    def load(cls: Type[T], path: Union[str, Path]) -> T:
+        return from_dict(cls, _load_file(path))
+
+    def save(self, path: Union[str, Path]):
+        path = Path(path)
+        if path.suffix in ('.yaml', '.yml'):
+            path.write_text(self.dumps_yaml())
+        else:
+            path.write_text(self.dumps_json())
+
+    def dumps_json(self, indent: int = 2) -> str:
+        return json.dumps(self.to_dict(), indent=indent)
+
+    def dumps_yaml(self) -> str:
+        import yaml
+        return yaml.safe_dump(self.to_dict(), sort_keys=False)

chugr/common/task_config.py

@@ -0,0 +1,24 @@
+from dataclasses import dataclass, field
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
+
+@dataclass
+class DataTaskCfg:
+    """
+    Attributes:
+        ...
+        output_tuple: Output tuples instead of dicts from task pipeline.
+        filter_valid: Filter out invalid / incomplete samples.
+        flatten_json: Flatten json dicts into parent sample dict. Common to have in wds datasets.
+        error_handler: 'specifies which error (exception) handler should be used, 'ignore_and_continue'
+            is good for training, 'reraise_exception' for debugging purposes.
+    """
+    decode_fn: Optional[Callable] = None
+    image_process_fn: Optional[Callable] = None
+    image_process_factory: Optional[Callable] = None  # seq_len -> image process fn (NaFlex bucketed batching)
+    text_process_fn: Optional[Callable] = None
+    decode_and_process_fn: Optional[Callable] = None
+    output_tuple: bool = False  # output features as tuple instead of dictionary
+    filter_valid: bool = False  # enable filter to keep samples with valid key-values
+    flatten_json: bool = True  # flatten nested 'json' dicts into parent sample
+    error_handler: str = 'reraise_exception'