diff --git a/scripts/create_testing_assets.py b/scripts/create_testing_assets.py deleted file mode 100644 index 68285ae..0000000 --- a/scripts/create_testing_assets.py +++ /dev/null @@ -1,20 +0,0 @@ -""" -Quick script to generate the assets for testing. -""" - -from __future__ import annotations - -import re -from pathlib import Path - -import unipercept as up - - -class TestingAssetDataset(up.data.sets.FolderPatternDataset, root=Path(__file__).parent / "assets" / "testing", pattern=re.compile(r"(\d{4})/(\d{6}).png$")): - Path(__file__).parent / "assets" / "testing", - re.compile(r"(\d{4})/(\d{6}).png$"), - depth_path=lambda p: - ) - -if __name__ == "__main__": - \ No newline at end of file diff --git a/sources/unipercept/cli/run.py b/sources/unipercept/cli/run.py new file mode 100644 index 0000000..95d36b1 --- /dev/null +++ b/sources/unipercept/cli/run.py @@ -0,0 +1,132 @@ +""" +Run a model in realtime or on a previously saved directory of images. +""" + +from __future__ import annotations + +import argparse +import os +import sys +import typing as T + +import torch +from omegaconf import DictConfig +from tabulate import tabulate + +import unipercept as up +from unipercept.cli._command import command + +_logger = up.log.get_logger() + + +KEY_SESSION_ID = "session_id" + + +@command(help="trian a model", description=__doc__) +@command.with_config +def train(p: argparse.ArgumentParser): + p_size = p.add_mutually_exclusive_group(required=False) + p_size.add_argument( + "--size", type=int, help="Size of the input images in pixels (smallest side)" + ) + p.add_argument( + "--weights", + "-w", + type=str, + help="path to load model weights from (overrides any state recovered by the config)", + ) + p.add_argument( + "--render", + type=str, + default="segmentation", + choices=["segmentation", "depth", "noop"], + help="rendering mode", + ) + + p.add_argument("input", type=str, default="0", help="input stream or directory") + + return _main + + +@torch.inference_mode() +def _main(args: argparse.Namespace): + config = args.config + + model = up.models.load(config.model) + preprocess = _build_transforms() + + if up.file_io.isdir(args.input): + run = _run_filesystem(model, preprocess, args.input) + else: + cap_num = int(args.input) + cap = _get_capture(cap_num) + run = _run_realtime(model, preprocess, cap) + + for inp, out in run: + print(out) + + +def _build_transforms(args): + import torchvision.transforms.v2 as transforms + + tf = [] + if args.size: + tf.append(transforms.Resize(args.size)) + + return up.data.ops.TorchvisionOp(tf) + + +def _get_capture(cap_num): + import cv2 + + cap = cv2.VideoCapture(cap_num, cv2.CAP_V4L2) + cap.set(cv2.CAP_PROP_FRAME_WIDTH, 224) + cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 224) + cap.set(cv2.CAP_PROP_FPS, 1) + return cap + + +def _run_realtime(model, preprocess, cap): + import numpy as np + import torchvision.transforms.v2 as transforms + + device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu") + + frame_num = 0 + while True: + ret, img_np = cap.read() + if not ret: + break + + # BGR -> RGB + img = transforms.functional.to_tensor(img_np[..., [2, 1, 0]]) + inp = up.create_inputs(img, frame_offset=frame_num) + inp = preprocess(inp) + out = model(inp) + + yield inp, out + + frame_num += 1 + + +def _run_filesystem(model, preprocess, path): + root = up.file_io.Path(path) + root_paths = list(root.iterdir()) + + if all(p.is_dir() for p in root_paths): + for p in root_paths: + yield from _run_filesystem(model, preprocess, p) + elif all(p.name.endswith(".png") for p in root_paths): + for p in root_paths: + img = up.data.tensors.Image.read(p) + inp = up.create_inputs(img, frame_offset=0) + inp = preprocess(inp) + out = model(inp) + + yield inp, out + else: + msg = ( + f"Invalid directory structure: {str(path)!r}, expected directories (sequence) " + "of PNG images sortable by name!" + ) + raise ValueError(msg) diff --git a/sources/unipercept/data/_loader.py b/sources/unipercept/data/_loader.py index 6d100a6..27b5dfc 100644 --- a/sources/unipercept/data/_loader.py +++ b/sources/unipercept/data/_loader.py @@ -3,13 +3,11 @@ from __future__ import annotations import abc -import dataclasses import dataclasses as D import enum import functools import itertools import math -import multiprocessing as M import operator import typing as T import warnings @@ -47,41 +45,34 @@ _logger = get_logger(__name__) -def _suggest_workers(): - """ - Suggests the number of workers for the dataloader based on the number of available CPUs - """ - try: - return max(cpus_available() - (cpus_available() // 4), 1) - except Exception: - return M.cpu_count() // get_process_count() - - -DEFAULT_NUM_WORKERS = get_env( - int, - "UP_DATALOADER_WORKERS", - "UNIPERCEPT_DATALOADER_WORKERS", - default=_suggest_workers(), -) - -DEFAULT_PREFETCH_FACTOR = get_env( - int, - "UP_DATALOADER_PREFETCH_FACTOR", - "UNIPERCEPT_DATALOADER_PREFETCH_FACTOR", - default=2, -) - - -@dataclasses.dataclass(slots=True, frozen=True) +@D.dataclass(slots=True, frozen=True) class DataLoaderConfig: """ Configuration parameters passed to the PyTorch dataoader """ drop_last: bool = False - pin_memory: bool = True - num_workers: int = DEFAULT_NUM_WORKERS - prefetch_factor: int | None = DEFAULT_PREFETCH_FACTOR + pin_memory: bool = D.field( + default_factory=lambda: get_env( + bool, + "UP_DATALOADER_PIN_MEMORY", + default=False, + ) + ) + num_workers: int = D.field( + default_factory=lambda: get_env( + int, + "UP_DATALOADER_WORKERS", + default=max(cpus_available(), 16), + ) + ) + prefetch_factor: int | None = D.field( + default_factory=lambda: get_env( + int, + "UP_DATALOADER_PREFETCH_FACTOR", + default=2, + ) + ) persistent_workers: bool | None = False @@ -90,7 +81,7 @@ class DataLoaderConfig: ################## -@dataclasses.dataclass(slots=True, frozen=True) +@D.dataclass(slots=True, frozen=True) class DataLoaderFactory: """ Factory for creating dataloaders. @@ -115,8 +106,8 @@ class DataLoaderFactory: dataset: PerceptionDataset actions: T.Sequence[Op] sampler: SamplerFactory - config: DataLoaderConfig = dataclasses.field(default_factory=DataLoaderConfig) - iterable: bool = dataclasses.field( + config: DataLoaderConfig = D.field(default_factory=DataLoaderConfig) + iterable: bool = D.field( default=False, metadata={ "help": ( @@ -142,11 +133,10 @@ def __call__( # Keyword arguments for the loader loader_kwargs = { - k: v for k, v in dataclasses.asdict(self.config).items() if v is not None + k: v for k, v in D.asdict(self.config).items() if v is not None } # Instantiate sampler - sampler_kwargs = {} if not use_distributed: sampler_kwargs["process_count"] = 1 @@ -187,7 +177,12 @@ def __call__( # Loader loader_kwargs["batch_size"] = batch_size loader_kwargs.setdefault("collate_fn", InputData.collate) - loader_kwargs.setdefault("worker_init_fn", _worker_init_fn) + + if loader_kwargs["num_workers"] > 0: + loader_kwargs.setdefault("worker_init_fn", _worker_init_fn) + elif "worker_init_fn" in loader_kwargs: + msg = f"Worker init function is set, but num_workers is {loader_kwargs['num_workers']}." + raise ValueError(msg) _logger.debug( "Creating dataloader (%d queued; %d × %d items):\n%s", @@ -683,7 +678,3 @@ def __init__( def __call__(self, queue: PerceptionDataqueue) -> Sampler: return self._fn(queue) - - -if __name__ == "__main__": - print("Default configuration for dataloader workers: ", DEFAULT_NUM_WORKERS) diff --git a/sources/unipercept/integrations/wandb_integration.py b/sources/unipercept/integrations/wandb_integration.py index 2c9daad..de44c6a 100644 --- a/sources/unipercept/integrations/wandb_integration.py +++ b/sources/unipercept/integrations/wandb_integration.py @@ -368,7 +368,10 @@ def artifact_historic_delete(artifact: wandb.Artifact, keep: int) -> None: api = wandb.Api() - vs = api.artifact_versions(type_name=artifact.type, name=name) + vs = [a.version for a in api.artifacts(type_name=artifact.type, name=name)] + if len(vs) <= keep: + _logger.info(f"Keeping all {name} artifacts") + return vs = sorted(vs, key=artifact_version_as_int, reverse=True) for artifact in vs[keep:]: _logger.info(f"Deleting artifact {name} version {artifact.version}") diff --git a/sources/unipercept/model.py b/sources/unipercept/model.py index 1f1befe..9dd4b08 100644 --- a/sources/unipercept/model.py +++ b/sources/unipercept/model.py @@ -346,7 +346,7 @@ def __call__( overrides_list = list(overrides) _logger.info( "Instantiating model with configuration overrides: %s", - ", ".join(overrides_list), + ", ".join(overrides_list) if len(overrides_list) > 0 else "(none)", ) model_config = apply_overrides(model_config, overrides_list) else: diff --git a/sources/unipercept/nn/layers/conv/_conditional.py b/sources/unipercept/nn/layers/conv/_conditional.py index 5ba5b48..c786404 100644 --- a/sources/unipercept/nn/layers/conv/_conditional.py +++ b/sources/unipercept/nn/layers/conv/_conditional.py @@ -17,7 +17,7 @@ from unipercept.utils.function import to_2tuple -from .utils import NormActivationMixin, PaddingMixin +from .utils import NormActivationMixin def get_condconv_initializer(initializer, num_experts, expert_shape): @@ -40,7 +40,7 @@ def condconv_initializer(weight): return condconv_initializer -class CondConv2d(PaddingMixin, NormActivationMixin, nn.Module): +class CondConv2d(NormActivationMixin, nn.Module): r""" Based on the implementation in `timm.layers`, where their docs state: > Conditionally Parameterized Convolution diff --git a/sources/unipercept/nn/layers/conv/_deform.py b/sources/unipercept/nn/layers/conv/_deform.py index b60173a..5105e87 100644 --- a/sources/unipercept/nn/layers/conv/_deform.py +++ b/sources/unipercept/nn/layers/conv/_deform.py @@ -20,7 +20,7 @@ from ..weight import init_xavier_fill_ from ._extended import Conv2d -from .utils import NormActivationMixin, PaddingMixin +from .utils import NormActivationMixin __all__ = ["ModDeform2d", "DeformConv2d"] @@ -223,7 +223,7 @@ def __repr__(self) -> str: return s -class ModDeform2d(NormActivationMixin, PaddingMixin, nn.Module): +class ModDeform2d(NormActivationMixin, nn.Module): """ Modulated deformable convolution. The offset mask is computed by a convolutional layer. diff --git a/sources/unipercept/nn/layers/conv/_extended.py b/sources/unipercept/nn/layers/conv/_extended.py index d64a8b4..d2c0928 100644 --- a/sources/unipercept/nn/layers/conv/_extended.py +++ b/sources/unipercept/nn/layers/conv/_extended.py @@ -7,23 +7,15 @@ import torch.nn as nn import typing_extensions as TX -from .utils import NormActivationMixin, PaddingMixin +from .utils import NormActivationMixin -__all__ = ["Conv2d", "Standard2d", "PadConv2d", "Separable2d"] +__all__ = ["Conv2d", "Standard2d", "Separable2d"] class Conv2d(NormActivationMixin, nn.Conv2d): pass -class PadConv2d(PaddingMixin, Conv2d): - @TX.override - def forward(self, x: torch.Tensor) -> torch.Tensor: - x = self._padding_forward(x, self.kernel_size, self.stride, self.dilation) - x = self._conv_forward(x, self.weight, self.bias) - return x - - class Standard2d(Conv2d): """ Implements weight standardization with learnable gain. @@ -43,7 +35,7 @@ def __init__(self, *args, gamma=1.0, eps=1e-6, gain=1.0, **kwargs): @TX.override def forward(self, x): - weight = F.batch_norm( + weight = nn.functional.batch_norm( self.weight.reshape(1, self.out_channels, -1), None, None, diff --git a/sources/unipercept/nn/layers/conv/utils.py b/sources/unipercept/nn/layers/conv/utils.py index 369c625..509455b 100644 --- a/sources/unipercept/nn/layers/conv/utils.py +++ b/sources/unipercept/nn/layers/conv/utils.py @@ -21,9 +21,6 @@ __all__ = [ "with_norm_activation", "NormActivationMixin", - "with_padding_support", - "Padding", - "PaddingMixin", ] @@ -67,7 +64,6 @@ def get_output_channels( # ---------------------- # -@torch.jit.ignore() def _init_sequential_norm_activation( cls, *args, @@ -85,7 +81,6 @@ def _init_sequential_norm_activation( return seq -@torch.jit.ignore() def _init_sequential_norm( cls, *args, norm: T.Optional[NormSpec], bias: T.Optional[bool] = None, **kwargs ) -> nn.Sequential: @@ -109,7 +104,6 @@ def _init_sequential_norm( return seq -@torch.jit.ignore() def _init_sequential_activation( cls, *args, activation: T.Optional[ActivationSpec] = None, **kwargs ) -> nn.Sequential: diff --git a/sources/unipercept/state.py b/sources/unipercept/state.py index e1b9d92..9915f2a 100644 --- a/sources/unipercept/state.py +++ b/sources/unipercept/state.py @@ -152,4 +152,10 @@ def gather_tensordict(td: TensorDictBase) -> TensorDict: def cpus_available(): - return len(os.sched_getaffinity(0)) + from multiprocessing import cpu_count + + try: + return len(os.sched_getaffinity(0)) + except: + pass + return max(cpu_count(), 1) diff --git a/tests/unipercept/conftest.py b/tests/unipercept/conftest.py index 928d85e..127a836 100644 --- a/tests/unipercept/conftest.py +++ b/tests/unipercept/conftest.py @@ -1,9 +1,8 @@ from __future__ import annotations +import os import re import typing as T -from dis import disco -from email.mime import image from pathlib import Path import pytest @@ -14,6 +13,10 @@ import unipercept as up from unipercept.data.sets import Metadata +os.environ["WANDB_MODE"] = "disabled" +os.environ["WANDB_SILENT"] = "true" +os.environ["UP_DATALOADER_WORKERS"] = "1" + ################################# # Directories with testing data # ################################# diff --git a/tests/unipercept/nn/layers/test_conv.py b/tests/unipercept/nn/layers/test_conv.py index a721064..38450b4 100644 --- a/tests/unipercept/nn/layers/test_conv.py +++ b/tests/unipercept/nn/layers/test_conv.py @@ -69,21 +69,7 @@ def test_forward( ) y = m.forward(input_tensor) - out = y.sum() - assert out.isfinite().all(), out - out.backward() - assert input_tensor.grad is not None - - print( - f"x : mean {input_tensor.mean().item(): 4.3f}, std {input_tensor.std().item(): 4.3f} | {tuple(input_tensor.shape)}" - ) - print( - f"y : mean {y.mean().item(): 4.3f}, std {y.std().item(): 4.3f} | {tuple(y.shape)}" - ) - print( - f"dx/dy : mean {input_tensor.grad.mean().item(): 4.3f}, std {input_tensor.grad.std().item(): 4.3f} | {tuple(input_tensor.grad.shape)}" - )