PTH: reading the tensors of a PyTorch checkpoint

Data Structures
struct	N_PTH
	A checkpoint's tensors. More...
struct	N_PTH_TENSOR
	One tensor of a checkpoint. More...

Macros
#define	N_PTH_MAX_DIMS   8
	Most dimensions a tensor read may have.

Functions
const N_PTH_TENSOR *	n_pth_find (const N_PTH *pth, const char *name)
	find a tensor by its path
void	n_pth_free (N_PTH **pth)
	release a checkpoint's tensors
N_PTH *	n_pth_load (const char *path)
	read every tensor of a checkpoint

Detailed Description

---

Data Structure Documentation

N_PTH

struct N_PTH

A checkpoint's tensors.

Examples

ex_pth.c.

Definition at line 80 of file n_pth.h.

Collaboration diagram for N_PTH:

Data Fields
size_t	count	how many
N_PTH_TENSOR *	tensors	the tensors, in the order the file lists them

N_PTH_TENSOR

struct N_PTH_TENSOR

One tensor of a checkpoint.

Examples

ex_pth.c.

Definition at line 66 of file n_pth.h.

Collaboration diagram for N_PTH_TENSOR:

Data Fields
size_t	count	how many
float *	data	its values, contiguous, row major; owned
int	dims	how many dimensions
char *	name	its path, the dictionary keys joined by dots; owned
int64_t	shape[8]	their sizes

Macro Definition Documentation

N_PTH_MAX_DIMS

#define N_PTH_MAX_DIMS   8

Most dimensions a tensor read may have.

Definition at line 63 of file n_pth.h.

Referenced by tensor_read().

Function Documentation

n_pth_find()

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

find a tensor by its path

Parameters

pth	the checkpoint
name	the path, "state_dict.backbone.0.0.weight"

Returns

the tensor, or NULL

Definition at line 776 of file n_pth.c.

References N_PTH::count, N_PTH_TENSOR::name, and N_PTH::tensors.

Referenced by main().

Here is the caller graph for this function:

n_pth_free()

void n_pth_free

(

N_PTH **

pth

)

release a checkpoint's tensors

Parameters

pth	the checkpoint, set to NULL

Definition at line 784 of file n_pth.c.

Referenced by main(), and n_pth_load().

Here is the caller graph for this function:

n_pth_load()

N_PTH * n_pth_load

(

const char *

path

)

read every tensor of a checkpoint

Parameters

path	the .pth or .pt file

Returns

the tensors, released with n_pth_free(), or NULL when the file is not a checkpoint this reads (not a ZIP, no data.pkl, a pickle it cannot follow, a storage missing or of another type)

Definition at line 738 of file n_pth.c.

References PK_VM::all, PK_VM::all_n, collect(), PK_OBJ::items, LOG_ERR, PK_VM::memo, n_log, n_pth_free(), n_zip_find_ending(), n_zip_read(), pk_run(), PTH_PICKLE_MAX, PK_OBJ::s, and PK_VM::stack.

Referenced by main().

Here is the call graph for this function:

Here is the caller graph for this function: