Big Blue Ceiling

PyTorch Utilities

Two small tools for the everyday friction of working with PyTorch.

Eddie OffermannDesign and engineering

Role
Design and engineering
Areas
AI, Tools
TorchVersionSpecifier's Gradio interface, with Python and CUDA version inputs and a list of matching torch, torchvision and torchaudio versions.

Two problems come up again and again in machine learning work: getting a PyTorch install that matches the machine, and figuring out what's in a checkpoint someone handed over. We wrote a small tool for each.

TorchVersionSpecifier

PyTorch, TorchVision and Torchaudio releases are tied to specific Python and CUDA versions, and working out a matching set by hand is tedious and easy to get wrong across many environments. Give TorchVersionSpecifier a Python version, a CUDA version and optionally a build variant such as cpu, and it lists the compatible version triples, ready to paste into a pip install. It pulls compatibility data directly from the PyTorch release sources, so it keeps up with new releases. It runs from the command line, in an interactive mode, or as a Gradio app.

Six-step flow in TorchVersionSpecifier: take a Python version and optional CUDA version and build variant, fetch PyTorch's stable wheel index, parse each wheel filename into package, version, build variant and Python tag, keep the matching ones and remove duplicates, pair each torch version with the same torchaudio version and a torchvision version from torchvision's compatibility table, and print the triples.
The tool reads PyTorch's own wheel index and torchvision's compatibility table, so it stays current without a hand-kept list.

TorchModelAnalyzer

A .pth file can hold a bare state dict, an optimizer state, params, params_ema or a whole pickled model, and the only way to tell is to load it and look. TorchModelAnalyzer does the looking: it reports the top-level keys, parameter counts and tensor shapes for each, the module hierarchy for saved models, and a summary of what the file contains, all as JSON that can be saved and compared between models.

Two-column comparison of TorchModelAnalyzer output. For a dictionary checkpoint it lists top-level keys, parameter counts and shapes for state_dict, optimizer keys, params_ema or params, a value or type for other keys, and summary flags. For a whole pickled model it reports the total parameter count, the module hierarchy, and leaf layer names and count.
A checkpoint dictionary and a whole saved model get different reports, both written as JSON.

Links

Sources

  • Aesopic

    AI-assisted long-form fiction, structured from series concept down to the page.

    • AI
    • Tools

  • AuRA Avatar Generator

    Customizable, photorealistic presenters for AI-generated video briefings.

    • AI
    • Tools
    • VFX

  • b2c3d

    Traditional 3D content creation, rebuilt with AI inside the tools.

    • AI
    • VFX
    • Pipelines
    • Tools

  • B2CQuantizer

    Batch quantization of Mistral models from a browser, on a rented GPU.

    • AI
    • Tools
    • Pipelines