Converting StyleGan to tensorflow js

General Discussion
, , ,
#1

Hello. This doubt may seem a little bit less researched, but here I go. I wanted to convert the stylegan generator to tensorflow js model. But this I understand that stylegan python (using TF2.x) has custom loops and objects in the generator object. I followed the steps mentioned below :

I right now have a keras model stylegan generator with weights loaded and the class has a call function that takes in noise and gives images. The definition is same as “StyleGAN_G” class in StyleGAN-Tensorflow2/stylegan.py at master · ialhashim/StyleGAN-Tensorflow2 (github.com). The weights I loaded was using the tutorial given in the same repository StyleGAN-Tensorflow2/stylegan.ipynb at master · ialhashim/StyleGAN-Tensorflow2 (github.com). So which format do I save the generator object (before converting it to tfjs model) ? Tensorflow seems to have a lot of save options

And after saving, what all flag values I have to set tfjsconvertor ?

1 Like
#2

@Jason might be able to help

1 Like
#3

Let me have a chat with the team and get back as I have not done much in the GAN space myself yet, but welcome to the community Saswat!

2 Likes
#4

Hi @Saswat_Das , please save it as SavedModel, model.save('my_model') (doc: Save and load Keras models  |  TensorFlow Core). And then follow the tfjs-converter doc to convert the SavedModel to tfjs format: tfjs/tfjs-converter at master · tensorflow/tfjs · GitHub
Basic cmd:

tensorflowjs_converter \
    --input_format=tf_saved_model \
    --output_format=tfjs_graph_model \
    --signature_name=serving_default \
    --saved_model_tags=serve \
    model_path/saved_model \
    new_path/web_model
4 Likes
#5

@lina128 Thank you . I converted the model to shard files and model.json. Then I loaded the graph model.

When I tried to run the following code to generate the image, it throws an error .
const noise = tf.randomNormal([512]).expandDims(0); let generatedImage = await generator.executeAsync(noise)

The error is that: util_base.ts:154 Uncaught (in promise) Error: Error in depthwiseConv2d: number of input channels (8) must match the inChannels dimension in filter 512.
at Object.assert (util_base.ts:154)
at depthwiseConv2d_ (depthwise_conv2d.ts:99)
at Object.depthwiseConv2d__op (operation.ts:51)
at Object.executeOp (convolution_executor.ts:195)
at operation_executor.ts:69
at engine.ts:467
at Engine.scopedRun (engine.ts:478)
at Engine.tidy (engine.ts:465)
at Object.tidy (globals.ts:192)
at operation_executor.ts:68

1 Like
#6

I am guessing it is because of the fact that depthwise convolution in tensorflow doesn’t support ‘BCHW’ format as mentioned in the API reference. But in the actual python code the data_format specified in the depthwise convolution was ‘BCHW’. Can anyone suggest what could be the error ?

1 Like
#7

It’s possible. Please add a feature request to add NCHW format or considering contribute. Or change the python model to use NHWC.

3 Likes
#8

Yes thank you. I changed the model to NHWC and it worked like a charm.

2 Likes
#9

Can you please share how did you convert it to NHWC format?

1 Like