Can I obtain the weights or constants from tf.save_model.load()?

General Discussion
, ,
#1

Hi everyone. I’m currently facing a challenge that I’m hoping to get some guidance on.

I had a model from tf.save_model.save(), and my goal is to load this model and extract the weights and constants contained within it. However, the loaded model is “_UserObject” and I cannot access the values inside. I’m curious if anyone has encountered a similar situation and found a solution or workaround.

Code is attached following. How can I get the parameters (weights/constants) from the model 2?

import tensorflow as tf

class Model(tf.Model):
    def __init__(self):
        super().__init__()
        self.const = tf.constant(123.456, dtype=tf.float32)

    @tf.function(jit_compile=True)
    def call(self, x):
        return self.const + x


input = tf.random.uniform([1, 2], dtype=tf.float32)

model1 = Model()

print(model1(input))

tf.saved_model.save(model1, "./aaa")

model2 = tf.saved_model.load("./aaa")

print(model2(input))

# how can I get the parameters (weights/constants) from the model 2?
#3

Hi @chijinz, Once you have loaded the model, you can use loaded_model.trainable_variables to get the model weights. In your case it would be

model2.trainable_variables

Thank You.

#5

@Kiran_Sai_Ramineni Thank you so much for your reply! It seems that trainable_variables only outputs weights that can be trained. It cannot find constant values. In my case, I would like to obtain the value of self.const, but this function cannot help. Code:

import tensorflow as tf

class Model(tf.keras.Model):
    def __init__(self):
        super().__init__()
        self.const = tf.constant(123.456, dtype=tf.float32)

    @tf.function(jit_compile=True)
    def call(self, x):
        return self.const + x


input = tf.random.uniform([1, 2], dtype=tf.float32)

model1 = Model()

print(model1(input))

tf.saved_model.save(model1, "./aaa")

model2 = tf.saved_model.load("./aaa")

print(model2(input))

print(model2.trainable_variables)

the output:

tf.Tensor([[124.317764 123.89178 ]], shape=(1, 2), dtype=float32)
tf.Tensor([[124.317764 123.89178 ]], shape=(1, 2), dtype=float32)
ListWrapper([])
#7

Hi everyone. I’m currently facing a challenge that I’m hoping to get some guidance on.

I had a model from tf.save_model.save(), and my goal is to load this model and extract the weights and constants contained within it. However, the loaded model is “_UserObject” and I cannot access the values inside. I’m curious if anyone has encountered a similar situation and found a solution or workaround.

Code:

import tensorflow as tf

class Model(tf.Model):
    def __init__(self):
        super().__init__()
        self.const = tf.constant(123.456, dtype=tf.float32)

    @tf.function(jit_compile=True)
    def call(self, x):
        return self.const + x


input = tf.random.uniform([1, 2], dtype=tf.float32)

model1 = Model()

print(model1(input))

tf.saved_model.save(model1, "./aaa")

model2 = tf.saved_model.load("./aaa")

print(model2(input))

# how can I get the parameters (weights/constants) from the model 2?
#8

Hi @chijinz, I tried with the alternative workaround to get the constants

import tensorflow as tf

class Model(tf.keras.Model):
    def __init__(self):
        super().__init__()
        self.const = tf.constant(123.456, dtype=tf.float32)

    @tf.function(jit_compile=True)
    def call(self, x):
        return self.const + x


input = tf.random.uniform([1, 2], dtype=tf.float32)

model1 = Model()

print(model1(input))

model1.save('my_model.keras')

new_model = tf.keras.models.load_model('my_model.keras',custom_objects={'Model': Model})

new_model.const.numpy() #123.456

Thank you.

#10

@Kiran_Sai_Ramineni Thank you again for your detailed reply! Perhaps I did not make my question clear. I cannot change the way of how the model is stored.

Suppose I have a model from tf.save_model.load(), how can I extract the constants contained within it? I highlight the part I cannot modify/observe.

import tensorflow as tf

input = tf.random.uniform([1, 2], dtype=tf.float32)

# =========
class Model(tf.Module):
    def __init__(self):
        super().__init__()
        self.const = tf.constant(123.456, dtype=tf.float32)

    @tf.function(jit_compile=True)
    def __call__(self, x):
        return self.const + x

model1 = Model()
tf.saved_model.save(model1, "./aaa")
# ======== what above is the part I cannot observe

model2 = tf.saved_model.load("./aaa")
print(model2(input))

# how can I get the parameters (weights/constants) from the model 2?