spandrel.util

package documentation
(source)

A module containing commonly-used functionality to implement architectures.

Package timm No package docstring; 4/5 functions, 2/2 classes, 2/3 modules documented

From __init__.py:

Class KeyCondition A condition that checks if a state dict has the given keys.
Function get_first_seq_index Returns the maximum index i such that key_pattern.format(str(i)) is in state.
Function get_pixelshuffle_params This will detect the upscale factor and number of features of a pixelshuffle module in the state dict.
Function get_scale_and_output_channels Returns a scale and number of output channels such that scale**2 * out_nc = x.
Function get_seq_len Returns the length of a sequence in the state dict.
Function store_hyperparameters Stores the hyperparameters of a class in a hyperparameters attribute.
def get_first_seq_index(state_dict: Mapping[str, object], key_pattern: str) -> int: (source)

Returns the maximum index i such that key_pattern.format(str(i)) is in state.

If no such key is in state, then -1 is returned.

Example:
get_first_seq_index(state, "body.{}.weight") -> -1 get_first_seq_index(state, "body.{}.weight") -> 3
def get_seq_len(state_dict: Mapping[str, object], seq_key: str) -> int: (source)

Returns the length of a sequence in the state dict.

The length is detected by finding the maximum index i such that {seq_key}.{i}.{suffix} is in state for some suffix.

Example:
get_seq_len(state, "body") -> 5
def get_scale_and_output_channels(x: int, input_channels: int) -> tuple[int, int]: (source)

Returns a scale and number of output channels such that scale**2 * out_nc = x.

This is commonly used for pixelshuffel layers.

def get_pixelshuffle_params(state_dict: Mapping[str, object], upsample_key: str = 'upsample', default_nf: int = 64) -> tuple[int, int]: (source)

This will detect the upscale factor and number of features of a pixelshuffle module in the state dict.

A pixelshuffle module is a sequence of alternating up convolutions and pixelshuffle. The class of this module is commonyl called Upsample. Examples of such modules can be found in most SISR architectures, such as SwinIR, HAT, RGT, and many more.

def store_hyperparameters(*, extra_parameters: Mapping[str, object] = {}): (source)

Stores the hyperparameters of a class in a hyperparameters attribute.