base

ModelSignatureWrapper(model)

Bases: Module

Model wrapper used to retrieve input shape. It can be used as a decorator of nn.Module, the first call to the forward method will retrieve the input shape and store it in the input_shapes attribute. It will also save the model summary in a file called model_summary.txt in the current working directory.

Source code in quadra/models/base.py

def __init__(self, model: nn.Module):
    super().__init__()
    self.instance = model
    self.input_shapes: Any = None
    self.disable = False

    if isinstance(self.instance, ModelSignatureWrapper):
        # Handle nested ModelSignatureWrapper
        self.input_shapes = self.instance.input_shapes
        self.instance = self.instance.instance

cpu(*args, **kwargs)

Handle calls to to method returning the underlying model.

Source code in quadra/models/base.py

def cpu(self, *args, **kwargs):
    """Handle calls to to method returning the underlying model."""
    self.instance = self.instance.cpu(*args, **kwargs)

    return self

forward(*args, **kwargs)

Retrieve the input shape and forward the model, if the input shape is already retrieved it will just forward the model.

Source code in quadra/models/base.py

def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor:
    """Retrieve the input shape and forward the model, if the input shape is already retrieved it will just forward
    the model.
    """
    if self.input_shapes is None and not self.disable:
        try:
            self.input_shapes = self._get_input_shapes(*args, **kwargs)
        except Exception:
            log.warning(
                "Failed to retrieve input shapes after forward! To export the model you'll need to "
                "provide the input shapes manually setting the config.export.input_shapes parameter! "
                "Alternatively you could try to use a forward with supported input types (and their compositions) "
                "(list, tuple, dict, tensors)."
            )
            self.disable = True

    return self.instance.forward(*args, **kwargs)

half(*args, **kwargs)

Handle calls to to method returning the underlying model.

Source code in quadra/models/base.py

def half(self, *args, **kwargs):
    """Handle calls to to method returning the underlying model."""
    self.instance = self.instance.half(*args, **kwargs)

    return self

to(*args, **kwargs)

Handle calls to to method returning the underlying model.

Source code in quadra/models/base.py

def to(self, *args, **kwargs):
    """Handle calls to to method returning the underlying model."""
    self.instance = self.instance.to(*args, **kwargs)

    return self