How to write numpy array to checkpoint in PyTorch

To write checkpoint in TensorFlow, first it will build the computational graph, the nodes and operations and how they are connected to each other and then it will use session.

Following codes storing the checkpoint in TensorFlow:

import os
import tensorflow.compat.v1 as tf
import numpy as np

def savecheckpoint(var_name, db, dtype, ckpt_path):
  with tf.Graph().as_default():
    init_value = tf.py_func(lambda: db, [], dtype, stateful=False)
    init_value.set_shape(db.shape)
    tf_db = tf.get_variable(var_name, initializer=init_value)
    saver = tf.train.Saver([tf_db])
    with tf.Session() as session:
      session.run(tf.global_variables_initializer())
      saver.save(session, ckpt_path)

path = 'home/Desktop/checkpoints/'
temp = ['1','2','3']
savecheckpoint("coref", np.array([m[0] for m in temp], dtype=np.int32), tf.int32, os.path.join(path, "var.npz"))

I am new to PyTorch, I have tried to convert above mentioed code in PyTorch:

import os
import torch
import numpy as np

def write_to_checkpoint(var_name, np_db, dtype, checkpoint_path):
    """Write np array to checkpoint."""
    t_db = torch.from_numpy(np_db)
    checkpoint = {
        var_name:t_db
    }
    torch.save(checkpoint,checkpoint_path+var_name+'.pth')

Are these codes equivalent? If not then please suggest correct equivalent codes.

I don’t know what the TF code is exactly doing, but in PyTorch you can save and load numpy arrays as well and adding them to a dict is also a valid idea:

a= np.random.randn(10)
checkpoint = {
    'np_arr': a
}
torch.save(checkpoint, 'tmp.pth')

b = torch.load('tmp.pth')['np_arr']

print(all(a == b))
> True
1 Like