cu.models: loading models

Load fastai and PyTorch models into a plain torch.nn.Module.

colabutilscs4337.models.load_model(path, arch=None, *, device=None, weights_only=None, strict=True, learner=False, eval=True)

Load a model file and return a torch.nn.Module (in eval mode by default).

Supported files:

• fastai learn.export() .pkl (returns learn.model, or the Learner with learner=True; needs the same custom functions/classes defined as when exported)

• fastai learn.save() .pth and PyTorch state dicts (need arch)

• full pickled PyTorch models (torch.save(model, ...))

• TorchScript (torch.jit.save) and torch.export .pt2 archives

arch is a model instance (for example torchvision.models.resnet18(num_classes=37)), a fastai Learner, or a zero-argument callable that builds one. device defaults to CUDA when available.

Parameters:

• path (str | PathLike)

• arch (Any | Callable[[], Any] | None)

• device (str | None)

• weights_only (bool | None)

• strict (bool)

• learner (bool)

• eval (bool)

colabutilscs4337.models.detect_format(path)

Guess the file format: fastai_export, torchscript, torch_export, or torch.

Parameters:

path (str | PathLike)

Return type:

str

colabutilscs4337.models.extract_state_dict(obj)

Pull the weights out of common checkpoint layouts.

Handles fastai learn.save ({"model": ..., "opt": ...}), Lightning-style {"state_dict": ...}, {"model_state_dict": ...}, and bare state dicts. A module. prefix from nn.DataParallel is removed.

Parameters:

obj (Mapping)

Return type:

dict