6 n__pth_8h
gullradriel edited this page 2026-10-09 00:55:18 +02:00

include/nilorea/n_pth.h

Reading the tensors of a PyTorch checkpoint (.pth, .pt) without Python. More...

Structures

Name
struct N_PTH_TENSOR
One tensor of a checkpoint.
struct N_PTH
A checkpoint's tensors.

Detailed Description

Reading the tensors of a PyTorch checkpoint (.pth, .pt) without Python.

Author: Castagnier Mickael

Version: 1.0

Date: 29/09/2026

A checkpoint saved by torch.save() since PyTorch 1.6 is a ZIP archive: a pickle, data.pkl, that describes the saved object, and one file a storage holding the raw numbers. This reads the pickle with a small interpreter that builds data only: dictionaries, lists, tuples, numbers, strings, and the three constructors a state dictionary uses (collections.OrderedDict, torch._utils._rebuild_tensor_v2 and _rebuild_parameter). Nothing named in the file is ever called; any other constructor becomes an opaque value, so a checkpoint's extra metadata does not stop its tensors being read.

Every tensor reachable through dictionaries is returned by its path, the keys joined by dots ("state_dict.backbone.0.0.weight"), as floats: float32, float16, bfloat16 and float64 storages are read and converted. Strided views are gathered. Little endian storages only, which is what every platform PyTorch runs on writes.

Every length, index and count in the file is checked against the data actually there: a damaged or hostile file is refused, not trusted.

Source code

/*
 * Nilorea Library
 * Copyright (C) 2005-2026 Castagnier Mickael
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
 * implied. See the License for the specific language governing
 * permissions and limitations under the License.
 *
 * SPDX-License-Identifier: Apache-2.0
 */


#ifndef __NILOREA_PTH__
#define __NILOREA_PTH__

#ifdef __cplusplus
extern "C" {
#endif


#include <stddef.h>
#include <stdint.h>

#define N_PTH_MAX_DIMS 8

typedef struct N_PTH_TENSOR {
    char* name;
    int dims;
    int64_t shape[N_PTH_MAX_DIMS];
    float* data;
    size_t count;
} N_PTH_TENSOR;

typedef struct N_PTH {
    N_PTH_TENSOR* tensors;
    size_t count;
} N_PTH;

N_PTH* n_pth_load(const char* path);

const N_PTH_TENSOR* n_pth_find(const N_PTH* pth, const char* name);

void n_pth_free(N_PTH** pth);


#ifdef __cplusplus
}
#endif

#endif