Any TensorFlow implementation?

[sayakpaul]

Hi.

This paper is simply great!

I was actually looking for a TensorFlow implementation, specifically TensorFlow-compatible checkpoints/frozengraphs/SavedModel. My plan is to convert it to a TensorFlow Lite model and create a demo mobile application it.

Looking forward to hearing from you.

[FriedRonaldo]

Unfortunately, we do not have any plan to implement TUNIT in tensorflow. However, you might refer to these articles for transferring the weights trained on pytorch to tensorflow; https://towardsdatascience.com/converting-a-simple-deep-learning-model-from-pytorch-to-tensorflow-b6b353351f5d https://forums.fast.ai/t/load-trained-pytorch-weights-in-tensorflow/11038 https://medium.com/@santhoshkolloju/convert-your-pytorch-models-to-tensorflow-with-onnx-84c3bdd8d722 I hope that these articles help you to port the pretrained weight (https://github.com/clovaai/tunit#pretrained-model).

[sayakpaul]

Thank you, @FriedRonaldo for the heads up. I will see what I can squeeze out of these. I will keep y'all posted here.

[sayakpaul]

Sorry about the reiteration here. Could you point me to a particular snippet that might be helpful for running inferences on a single pair of images? I think that way the ONNX conversion process would be a bit easier.

[FriedRonaldo]

Actually, I have no experience in executing models on mobile platforms. The articles are from google. If your goal is just to convert the weights into a form of tensorflow, there might be many GitHub open sources. I think that the way you mentioned; pytorch->ONNX->TF seems to be easier than other ways.

I tested the method in (https://towardsdatascience.com/converting-a-simple-deep-learning-model-from-pytorch-to-tensorflow-b6b353351f5d) with a simple model that takes two inputs (not TUNIT). I check that the export (torch.onnx.export ~) works, however, I do not test after this step because I do not have time enough to test all the porting procedure now. I am sorry for that. After my own work, I will try it with TUNIT.

I hope that it works for you.

[sayakpaul]

@FriedRonaldo thank you for being so generous. I was actually referring to the following workflow for running a demo inference:

• Load up the pre-trained checkpoints for a particular dataset.
• Take two images i.e. source and reference.
• Run inference with the pre-trained checkpoints loaded in the first step.

Could you point me to a snippet in the current codebase of tunit that could help me implement this? Let me know if anything is unclear.

[FriedRonaldo]

Oh, I am sorry. I misunderstood your intention.

Is it correct that you ask about the important code parts to conduct the process you mentioned? If it is right, it might be like below:

If your images are loaded by cv2.imread or PIL,

0. Read images in tensor
1. Build the model with build_model (https://github.com/clovaai/tunit/blob/master/main.py#L341)
2. Load weights by load_model (https://github.com/clovaai/tunit/blob/master/main.py#L404) to load pretrained net
3. Test with a loaded network.

args might contain 'gpu', 'distributed', 'batch_size', 'workers', 'ngpus_per_node', 'img_size', 'sty_dim', 'output_k', 'load_model', 'multiprocessing_distributed', 'start_epoch', 'log_dir'. (But you can skip many of them by modifying build_model and load_model If you run the inference only, start_epoch, 'multiprocessing_distributed', 'workers', 'ngpus_per_node' can be skipped by removing some "if" statements.)

import models and other
init. args

networks, _ = build_model(args)
load_model(args, networks, {})
C_EMA = networks['C_EMA']
G_EMA = networks['G_EMA']

Then generate images like 
from https://github.com/clovaai/tunit/blob/master/validation/validation.py#L143
to https://github.com/clovaai/tunit/blob/master/validation/validation.py#L145

x_src and x_ref are tensor shaped in [1, 3, 128, 128].

c_src = G_EMA.cnt_encoder(x_src)
s_ref = C_EMA(x_ref, sty=True)
x_res = G_EMA.decode(c_src, s_ref)

you can also modify the build_model and load_model to remove the discriminator and optimizers.

Does it correspond to your intention? I understand the comment as you want to write a simplified script to conduct an inference with a single image.

[FriedRonaldo]

Here, this is the simplest code:

import torch
from models.generator import Generator
from models.guidingNet import GuidingNet
import torch.nn.functional as F
import torchvision.utils as vutils
from PIL import Image
from torchvision.transforms import ToTensor

G = Generator(128, 128)
C = GuidingNet(128)

load_file = './logs/animalFaces10_0_00/model_4568.ckpt'

checkpoint = torch.load(load_file, map_location='cpu')
G.load_state_dict(checkpoint['G_EMA_state_dict'])
C.load_state_dict(checkpoint['C_EMA_state_dict'])

G.eval()
C.eval()

x_src = ToTensor()(Image.open('human1.png')).unsqueeze(0)
x_ref = ToTensor()(Image.open('source.jpg')).unsqueeze(0)

x_src = F.interpolate(x_src, size=(128, 128))
x_ref = F.interpolate(x_ref, size=(128, 128))

x_src = (x_src - 0.5) / 0.5
x_ref = (x_ref - 0.5) / 0.5

c_src = G.cnt_encoder(x_src)
s_ref = C.moco(x_ref)
x_res = G.decode(c_src, s_ref)

vutils.save_image(x_res, 'test_out.jpg', normalize=True, padding=0)

I write this after your comment. You should change x_src and x_ref to your exmaples. If you run this with noise vectors, it produces a weird one.

Totally random texture. (Somewhat scary) - gaussian source and reference.

Real example: reference source result

reference: source:

result:

[sayakpaul]

Thank you very much for your detailed guidance. I genuinely appreciate this. I will try putting together a Colab Notebook from this conversation and will also try including the TensorFlow conversion steps for the community to get benefitted.

Thank you once again :)

[FriedRonaldo]

Glad to be helpful!

[sayakpaul]

@FriedRonaldo I was able to put together a Colab Notebook from your code snippet. Currently, I am having to first copy over the checkpoint files to my personal Drive. This could have been a bit more accessible if you hosted the checkpoint files as a GitHub Release. It would allow the users to easily retrieve them using wget.

Regarding ONNX export, I have raised an issue on the PyTorch forums. I Will update here if I can make some progress.

[FriedRonaldo]

Ok. I got it. I am not used to GitHub Release, however, I will try it to make this convenient. I saw your Colab notebook and the result. It is good work! But there must be a constraint on the choice of the reference image. The reference image you used is not a domain included in the training.

The issue on ONNX seems that forward function is not implemented. I added the implementation of forward to solve the problem. The inference code changes from

c_src = G.cnt_encoder(x_src)
s_ref = C.moco(x_ref)
x_res = G.decode(c_src, s_ref)

to

s_ref = C.moco(x_ref)
x_res = G(x_src, s_ref)

---

Edit: After reading your colab code, I decided to write a code for inference with a single image in Colab. I will be uploaded in four days.

[sayakpaul]

I am not used to GitHub Release, however, I will try it to make this convenient.

So, here's a workflow that you could follow:

• Simply zip the individual folders and name them accordingly.
• Add the different zipped checkpoints to a release.
• Publish the release.

GitHub Releases also allow you to simply drag and drop files to added them to a particular release. So, it's really convenient.

But there must be a constraint on the choice of the reference image. The reference image you used is not a domain included in the training.

Sure, I can add a note on that in the Notebook itself.

After reading your colab code, I decided to write a code for inference with a single image in Colab.

Great! Do you think we could collaborate on that? We already have a decent Colab Notebook ready. In order to make the user experience a bit more seamless, you could create the GitHub release with the model checkpoints. Meanwhile, I can polish the notebook with more annotations, etc. We can then bounce back feedback. What do you think?

[FriedRonaldo]

So, here's a workflow that you could follow: Simply zip the individual folders and name them accordingly. Add the different zipped checkpoints to a release. Publish the release. GitHub Releases also allow you to simply drag and drop files to added them to a particular release. So, it's really convenient.

Thanks for your suggestion! I will try it with the workflow.

Great! Do you think we could collaborate on that? We already have a decent Colab Notebook ready. In order to make the user experience a bit more seamless, you could create the GitHub release with the model checkpoints. Meanwhile, I can polish the notebook with more annotations, etc. We can then bounce back feedback. What do you think?

I think it is a good idea. However, I am busy with my personal work till Saturday(KST 6/20), though, I will try keeping to communicate with you as fast as I can.

[sayakpaul]

Of course! I appreciate it. Meanwhile, I will add the necessary comments in the Notebook.

[FriedRonaldo]

Thanks for your effort!

[victorca25]

I am not used to GitHub Release, however, I will try it to make this convenient.

So, here's a workflow that you could follow:

• Simply zip the individual folders and name them accordingly.

• Add the different zipped checkpoints to a release.

• Publish the release.

You can also do a simple !wget to download the models from drive, no need to put them in the release

[FriedRonaldo]

@victorca25 Thanks for the suggestion. I will try wget with drive first.

[sayakpaul]

@victorca25 I don't think that's possible from the direct shareable link one would get from the .ckpt files. Here's an example: https://www.loom.com/share/d7a41d9981674d71a62ce3ce1f7dd685.

[victorca25]

@sayakpaul works for me

Edit: this may help https://stackoverflow.com/questions/25010369/wget-curl-large-file-from-google-drive

[sayakpaul]

Well, it comes with specifications if I understood correctly. But I don't think one would be able to use wget or curl to retrieve the checkpoints using their shareable links. Hosting them as GitHub release might be an easier solution.

I did create a Gist a couple of months back that demonstrates the entire process of retrieving Drive files (publically shared) using wget.

[FriedRonaldo]

@sayakpaul I added the file animalFaces10_0_00.ckpt to release as a form of .zip. Sorry for late.

[sayakpaul]

No worries, @FriedRonaldo. Here's the Colab Notebook, I prepared with your pointers: https://github.com/sayakpaul/Adventures-in-TensorFlow-Lite/blob/master/TUNIT_Conversion_to_TF_Lite.ipynb.

[FriedRonaldo]

@sayakpaul Good!