Using your own model or built-in models¶

Own model¶

RAITAP allows you to use your own model in any of the following supported formats:

.pt / .pth containing a state_dict — recommended. Combine with model.arch and model.num_classes to instantiate a torchvision architecture and load weights into it. Portable across environments and torchvision versions.
.pt / .pth containing a TorchScript archive (saved via torch.jit.save(scripted, path)). Self-contained — no arch / num_classes needed.
.pt / .pth containing a full pickled torch.nn.Module (saved via torch.save(model, path)). Deprecated and refused by default — this format requires unsafe pickle deserialisation, which executes arbitrary code embedded in the file. Opt in only for fully trusted files; see --allow-unsafe-pickle. The pickle also embeds fully-qualified class paths so it breaks when classes are renamed or when torchvision is bumped. Migrate with one line (in a trusted environment):
```
m = torch.load("model.pth", weights_only=False)
torch.save(m.state_dict(), "weights.pth")
```
.onnx
.ubj: XGBoost native binary (fitted XGBClassifier saved via estimator.save_model("model.ubj")). Requires the xgboost extra: uv sync --extra xgboost. Use with shap.TreeExplainer (also needs --extra shap) or any model-agnostic SHAP explainer.

Set the source option to the path of your model file (see Configuration). For state-dict loading also set arch and num_classes:

YAML

model:
  source: "weights.pth"
  arch: "resnet18"
  num_classes: 2

Python

from raitap.models import ModelConfig

model = ModelConfig(
    source="weights.pth",
    arch="resnet18",
    num_classes=2,
)

Not all modules and underlying libraries support all formats. See each module's library compatibility page for more details.

Built-in models¶

Alternatively, you can use any model provided by the torchvision.models library. They will be initialised with weights="DEFAULT". You can find the list here.

You simply need to set the source option to the name of the model (see Configuration).