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.

Welcome to the TF community forum.

As It Is more related to pytorch I suggest to post the same on:

1 Like

I have posted same question in PyTorch community too. Thanks