How to train/load S/M/L CNN models in tensorflow?

[kristoftunner]

Is there a way to train/load S/M/L CNN models in tensorflow? I am interested in experimenting a bit with these models in tensorflow or onnxruntime. I see that there is one specific model in the tensorflow directory, but I am not sure which one is it.

[Tama47]

Is there a way to train/load S/M/L CNN models in tensorflow?

Yes, you would need to load the original models in TensorFlow.

I see that there is one specific model in the tensorflow directory, but I am not sure which one is it.

Someone has converted the original Anime4K models into Core ML models. I can provide you the link.

The ones you're looking for are under Models >
model-sr-s.wifm / model-sr-m.wifm / model-sr-l.wifm for upscale models
model-restore-s.wifm / model-restore-m.wifm / model-restore-l.wifm for restore models

You would need to convert them to TensorFlow, then create a Python or Jupyter Notebook script to load the weights and models. You can use the models to fine-tune and train your own, better model.

Note

I have not converted or trained the models myself, and cannot guarantee success. I can only provide general steps, and you will need to do your own research. Supposedly, the steps to convert between Core ML and TensorFlow should be relatively straightforward. The training process itself should be more or less the same as training any other TensorFlow or ESRGAN models.

Sample Python Script:

import coremltools as ct



# Load Core ML model
coreml_model = ct.models.MLModel('model-sr-s.wifm')



# Convert Core ML to TensorFlow
tf_model = ct.convert(coreml_model, source='mlmodel', target='tensorflow')

[arianaa30]

@Tama47 The training code located in \tensorflow dir is for the restore or upscale model? And if it is the restore, is it easy to change it to the "upscale" model to train?

[Fannovel16]

@Tama47 From what I've researched so far, there is no way to convert current version of MLModel to TF2 or ONNX. However, I managed to get Netron working and also loading weight:

1. Change the file extension from .whml to .zip
2. Compress all files inside of *.mlpackage folder (not including the folder itself) to a zip file

Preset-a-hq:

[Fannovel16]

I can convert some GLSL files to PyTorch now but still stuck at converting the weight. Here is the code if anyone interested:
https://colab.research.google.com/drive/11xAn4fyAUJPZOjrxwnL2ipl_1DGGegkB

[Fannovel16]

@arianaa30

The training code located in \tensorflow dir is for the restore or upscale model?

It contains both. In Gen_Shader.ipynb, SR1Model generate restore models while SR2Model generate upscale ones

[arianaa30]

@arianaa30

The training code located in \tensorflow dir is for the restore or upscale model?

It contains both. In Gen_Shader.ipynb, SR1Model generate restore models while SR2Model generate upscale ones

Thanks. Apparently the models there for for M size shaders. Do you happen to know what parameter values (block_depth, etc?) I should use to get the S/L sizes?

Also, the code uses epochs=1 (3 times). Should I change them to like 100? I noticed the loss doesn't really decrease.

[Fannovel16]

@arianaa30

Thanks. Apparently the models there for for M size shaders. Do you happen to know what parameter values (block_depth, etc?) I should use to get the S/L sizes?

Ig you can figure out the block_depth with a model's components
Conv2d(3, 4) means it has 3 input channels and 4 output channels. CReLU() activation function doubles channel size, e.g. (1, 128, 128, 4) -> (1, 128, 128, 8)

Size S:

Anime4K(
  (conv2d_0): Sequential(
    (0): Conv2d(3, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): CReLU()
  )
  (conv2d_1_0): Sequential(
    (0): Conv2d(8, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): CReLU()
  )
  (conv2d_2_0): Sequential(
    (0): Conv2d(8, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): CReLU()
  )
  (conv2d_last_0): Conv2d(8, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (upsample): Upsample(scale_factor=2.0, mode='bilinear')
  (depth_to_space): PixelShuffle(upscale_factor=2)
)

Size M

Anime4K(
  (conv2d_0): Sequential(
    (0): Conv2d(3, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): CReLU()
  )
  (conv2d_1_0): Sequential(
    (0): Conv2d(8, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): CReLU()
  )
  (conv2d_2_0): Sequential(
    (0): Conv2d(8, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): CReLU()
  )
  (conv2d_3_0): Sequential(
    (0): Conv2d(8, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): CReLU()
  )
  (conv2d_4_0): Sequential(
    (0): Conv2d(8, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): CReLU()
  )
  (conv2d_5_0): Sequential(
    (0): Conv2d(8, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): CReLU()
  )
  (conv2d_6_0): Sequential(
    (0): Conv2d(8, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): CReLU()
  )
  (conv2d_last_0): Conv2d(56, 4, kernel_size=(1, 1), stride=(1, 1))
  (upsample): Upsample(scale_factor=2.0, mode='bilinear')
  (depth_to_space): PixelShuffle(upscale_factor=2)
)

Size L:

Anime4K(
  (conv2d_0): Sequential(
    (0): Conv2d(3, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): CReLU()
  )
  (conv2d_1): Sequential(
    (0): Conv2d(3, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): CReLU()
  )
  (conv2d_1_0): Sequential(
    (0): Conv2d(16, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): CReLU()
  )
  (conv2d_1_1): Sequential(
    (0): Conv2d(16, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): CReLU()
  )
  (conv2d_2_0): Sequential(
    (0): Conv2d(16, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): CReLU()
  )
  (conv2d_2_1): Sequential(
    (0): Conv2d(16, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): CReLU()
  )
  (conv2d_last_0): Conv2d(16, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (conv2d_last_1): Conv2d(16, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (conv2d_last_2): Conv2d(16, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (upsample): Upsample(scale_factor=2.0, mode='bilinear')
  (depth_to_space): PixelShuffle(upscale_factor=2)
)

Size VL:

Anime4K(
  (conv2d_0): Sequential(
    (0): Conv2d(3, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): CReLU()
  )
  (conv2d_1): Sequential(
    (0): Conv2d(3, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): CReLU()
  )
  (conv2d_1_0): Sequential(
    (0): Conv2d(16, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): CReLU()
  )
  (conv2d_1_1): Sequential(
    (0): Conv2d(16, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): CReLU()
  )
  (conv2d_2_0): Sequential(
    (0): Conv2d(16, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): CReLU()
  )
  (conv2d_2_1): Sequential(
    (0): Conv2d(16, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): CReLU()
  )
  (conv2d_3_0): Sequential(
    (0): Conv2d(16, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): CReLU()
  )
  (conv2d_3_1): Sequential(
    (0): Conv2d(16, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): CReLU()
  )
  (conv2d_4_0): Sequential(
    (0): Conv2d(16, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): CReLU()
  )
  (conv2d_4_1): Sequential(
    (0): Conv2d(16, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): CReLU()
  )
  (conv2d_5_0): Sequential(
    (0): Conv2d(16, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): CReLU()
  )
  (conv2d_5_1): Sequential(
    (0): Conv2d(16, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): CReLU()
  )
  (conv2d_6_0): Sequential(
    (0): Conv2d(16, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): CReLU()
  )
  (conv2d_6_1): Sequential(
    (0): Conv2d(16, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): CReLU()
  )
  (conv2d_last_0): Conv2d(112, 4, kernel_size=(1, 1), stride=(1, 1))
  (conv2d_last_1): Conv2d(112, 4, kernel_size=(1, 1), stride=(1, 1))
  (conv2d_last_2): Conv2d(112, 4, kernel_size=(1, 1), stride=(1, 1))
  (upsample): Upsample(scale_factor=2.0, mode='bilinear')
  (depth_to_space): PixelShuffle(upscale_factor=2)
)

Size UL:

Anime4K(
  (conv2d_0): Sequential(
    (0): Conv2d(3, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): CReLU()
  )
  (conv2d_1): Sequential(
    (0): Conv2d(3, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): CReLU()
  )
  (conv2d_2): Sequential(
    (0): Conv2d(3, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): CReLU()
  )
  (conv2d_1_0): Sequential(
    (0): Conv2d(24, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): CReLU()
  )
  (conv2d_1_1): Sequential(
    (0): Conv2d(24, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): CReLU()
  )
  (conv2d_1_2): Sequential(
    (0): Conv2d(24, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): CReLU()
  )
  (conv2d_2_0): Sequential(
    (0): Conv2d(24, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): CReLU()
  )
  (conv2d_2_1): Sequential(
    (0): Conv2d(24, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): CReLU()
  )
  (conv2d_2_2): Sequential(
    (0): Conv2d(24, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): CReLU()
  )
  (conv2d_3_0): Sequential(
    (0): Conv2d(24, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): CReLU()
  )
  (conv2d_3_1): Sequential(
    (0): Conv2d(24, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): CReLU()
  )
  (conv2d_3_2): Sequential(
    (0): Conv2d(24, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): CReLU()
  )
  (conv2d_4_0): Sequential(
    (0): Conv2d(24, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): CReLU()
  )
  (conv2d_4_1): Sequential(
    (0): Conv2d(24, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): CReLU()
  )
  (conv2d_4_2): Sequential(
    (0): Conv2d(24, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): CReLU()
  )
  (conv2d_5_0): Sequential(
    (0): Conv2d(24, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): CReLU()
  )
  (conv2d_5_1): Sequential(
    (0): Conv2d(24, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): CReLU()
  )
  (conv2d_5_2): Sequential(
    (0): Conv2d(24, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): CReLU()
  )
  (conv2d_6_0): Sequential(
    (0): Conv2d(24, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): CReLU()
  )
  (conv2d_6_1): Sequential(
    (0): Conv2d(24, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): CReLU()
  )
  (conv2d_6_2): Sequential(
    (0): Conv2d(24, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): CReLU()
  )
  (conv2d_last_0): Conv2d(120, 4, kernel_size=(1, 1), stride=(1, 1))
  (conv2d_last_1): Conv2d(120, 4, kernel_size=(1, 1), stride=(1, 1))
  (conv2d_last_2): Conv2d(120, 4, kernel_size=(1, 1), stride=(1, 1))
  (upsample): Upsample(scale_factor=2.0, mode='bilinear')
  (depth_to_space): PixelShuffle(upscale_factor=2)
)

[arianaa30]

Thanks. I have those architecture. But do you know what to pass to this function to get each of those S, L, VL sizes? I need it for training.

def SR2Model (input_depth=3, highway_depth=4, block_depth=4, init='he_normal', init_last = RandomNormal (mean=0.0, stddev=0.001)):

[Fannovel16]

@arianaa30 My main library is PyTorch so Idk tbh

[arianaa30]

@Fannovel16 Btw, do you know how to measure SSIM/PSNR of what Anime4K shaders provide me (upscaled version of low-res image) vs the original high resolution image? Is there a way to measure them?

[Fannovel16]

@arianaa30 You can pass images to mpv, ffmpeg with compiled libplacebo using commands and save upscaled images. I'm not sure how to do that tho since I'm not familiar much with ffmpeg or mpv. Alternatively, you can try Anime4K-GPU and use canvas.toDataURL("image/png") to save the results

[Fannovel16]

Btw I have no luck on converting GLSL shaders into actual models in PyTorch. I believe the problem is transforming the weight.

[arianaa30]

@arianaa30 You can pass images to mpv, ffmpeg with compiled libplacebo using commands and save upscaled images. I'm not sure how to do that tho since I'm not familiar much with ffmpeg or mpv. Alternatively, you can try Anime4K-GPU and use canvas.toDataURL("image/png") to save the results

Hmm ok thanks. The problem is we apply multiple anime4k shaders (restore, upscale, restore, ...). Not sure if we can do that..