tensorflow / compression

Data compression in TensorFlow
Apache License 2.0
860 stars 251 forks source link

Unable to save model #174

Open mohit-av opened 1 year ago

mohit-av commented 1 year ago

I'm creating a model to decompress my image in TF using Conv2DTranspose and GDN layer. My architecture looks like this:

import os
import tensorflow as tf
from tensorflow_compression import layers
from tensorflow.keras.layers import Conv2DTranspose
from tensorflow.keras.models import Sequential
network_channels = 128
compression_channels = 192
in_channels = 3

gdn_layer_1 = layers.GDN(inverse=True, rectify=False, data_format = "channels_first") 
gdn_layer_2 = layers.GDN(inverse=True, rectify=False, data_format = "channels_first")
gdn_layer_3 = layers.GDN(inverse=True, rectify=False, data_format = "channels_first")
model = Sequential()
model.add(Conv2DTranspose(network_channels, (5, 5), strides = (2, 2), padding= 'same', output_padding = (1,1), data_format = 'channels_first', input_shape=( compression_channels, 1, 1)))
model.add(gdn_layer_1)
model.add(Conv2DTranspose(network_channels, (5, 5), strides = (2, 2), padding= 'same', output_padding = (1,1), data_format = 'channels_first'))
model.add(gdn_layer_2)
model.add(Conv2DTranspose(network_channels, (5, 5), strides = (2, 2), padding= 'same', output_padding = (1,1), data_format = 'channels_first'))
model.add(gdn_layer_3)
model.add(Conv2DTranspose(in_channels, (5, 5), strides = (2, 2), padding= 'same', output_padding = (1,1), data_format = 'channels_first'))
model.summary()

Screenshot 2023-02-22 at 9 55 13 AM

When I'm saving my model using the below command.

model.save(os.path.join("./tf_model"))

It's resulting in this issue:

Screenshot 2023-02-22 at 9 58 41 AM

My TF version is 2.11.0. I'm also attaching link to Colab Notebook to replicate the issue.