Weight Converter & Pretrained Loaders API¶

Auto-generated from source — see Weight Porting and Pretrained Weights for narrative guides.

Converter¶

keras_climate.weights.converter.WeightConverter ¶

WeightConverter(keras_model, torch_state_dict, name_map, param_kind_map=None, skip_patterns=None)

Source code in keras_climate/weights/converter.py

def __init__(
    self,
    keras_model,
    torch_state_dict,
    name_map,
    param_kind_map=None,
    skip_patterns=None,
):
    self.model = keras_model
    self.state_dict = torch_state_dict
    self.name_map = name_map if callable(name_map) else build_regex_mapper(name_map)
    self.param_kind_map = param_kind_map or {}
    self.skip_patterns = [re.compile(p) for p in (skip_patterns or [])]

keras_climate.weights.converter.infer_transpose ¶

infer_transpose(torch_shape, keras_shape, param_kind=None)

Source code in keras_climate/weights/converter.py

def infer_transpose(torch_shape, keras_shape, param_kind=None):
    if param_kind == "dense_kernel":
        return (1, 0)

    if tuple(torch_shape) == tuple(keras_shape):
        return None

    if len(torch_shape) == 2 and torch_shape[::-1] == tuple(keras_shape):
        return (1, 0)

    if param_kind == "conv2d_kernel" or len(torch_shape) == 4:
        candidate = (torch_shape[2], torch_shape[3], torch_shape[1], torch_shape[0])
        if candidate == tuple(keras_shape):
            return (2, 3, 1, 0)

    if param_kind == "conv3d_kernel" or len(torch_shape) == 5:
        candidate = (
            torch_shape[2],
            torch_shape[3],
            torch_shape[4],
            torch_shape[1],
            torch_shape[0],
        )
        if candidate == tuple(keras_shape):
            return (2, 3, 4, 1, 0)

    if param_kind == "conv1d_kernel" or len(torch_shape) == 3:
        candidate = (torch_shape[2], torch_shape[1], torch_shape[0])
        if candidate == tuple(keras_shape):
            return (2, 1, 0)

    if param_kind == "depthwise_kernel" and len(torch_shape) == 4:
        return "depthwise"

    return None

keras_climate.weights.converter.apply_transpose ¶

apply_transpose(array, perm)

Source code in keras_climate/weights/converter.py

def apply_transpose(array, perm):
    if perm is None:
        return array
    if perm == "depthwise":
        out_ch, _, kh, kw = array.shape
        arr = array.reshape(out_ch, 1, kh, kw)
        arr = np.transpose(arr, (2, 3, 0, 1))
        return arr
    return np.transpose(array, perm)

keras_climate.weights.converter.build_regex_mapper ¶

build_regex_mapper(rules)

Source code in keras_climate/weights/converter.py

def build_regex_mapper(rules):

    compiled = [(re.compile(p), r) for p, r in rules]

    def mapper(torch_key):
        for pattern, replacement in compiled:
            if pattern.search(torch_key):
                return pattern.sub(replacement, torch_key)
        return None

    return mapper

keras_climate.weights.converter.load_torch_state_dict_as_numpy ¶

load_torch_state_dict_as_numpy(checkpoint_path, key_prefix_strip=None)

Source code in keras_climate/weights/converter.py

def load_torch_state_dict_as_numpy(checkpoint_path, key_prefix_strip=None):
    import torch

    state_dict = torch.load(checkpoint_path, map_location="cpu")
    if "state_dict" in state_dict:
        state_dict = state_dict["state_dict"]
    if "model" in state_dict and all(
        isinstance(v, dict) for v in [state_dict.get("model", {})]
    ):
        state_dict = state_dict.get("model", state_dict)

    out = {}
    for k, v in state_dict.items():
        if key_prefix_strip and k.startswith(key_prefix_strip):
            k = k[len(key_prefix_strip) :]
        out[k] = v.detach().cpu().numpy() if hasattr(v, "detach") else np.asarray(v)
    return out

keras_climate.weights.converter.load_safetensors_as_numpy ¶

load_safetensors_as_numpy(checkpoint_path)

Source code in keras_climate/weights/converter.py

def load_safetensors_as_numpy(checkpoint_path):
    from safetensors import safe_open

    out = {}
    with safe_open(checkpoint_path, framework="numpy") as f:
        for key in f.keys():
            out[key] = f.get_tensor(key)
    return out