<jupyter_start><jupyter_text>Example of weight loading Loading the weights in TensorflowThe link used is the one for the A5 model, and this example will continue loading the weights for that specific model. To load the weights of a different model change now the link of the model. A0 : "https://tfhub.dev/tensorflow/movinet/a0/base/kinetics-600/classification/2"A1 : "https://tfhub.dev/tensorflow/movinet/a1/base/kinetics-600/classification/2"A2 : "https://tfhub.dev/tensorflow/movinet/a2/base/kinetics-600/classification/2"A3 : "https://tfhub.dev/tensorflow/movinet/a3/base/kinetics-600/classification/2"A4 : "https://tfhub.dev/tensorflow/movinet/a4/base/kinetics-600/classification/2"A5 : "https://tfhub.dev/tensorflow/movinet/a5/base/kinetics-600/classification/2"<jupyter_code>inputs = tf.keras.layers.Input(
    shape=[None, None, None, 3],
    dtype=tf.float32)

encoder = hub.KerasLayer(
    "https://tfhub.dev/tensorflow/movinet/a0/base/kinetics-600/classification/3")

# Important: due to a bug in the tf.nn.conv3d CPU implementation, we must
# compile with tf.function to enforce correct behavior. Otherwise, the output
# on CPU may be incorrect.
encoder.call = tf.function(encoder.call, experimental_compile=True)

# [batch_size, 600]
outputs = encoder(dict(image=inputs))

model = tf.keras.Model(inputs, outputs)

#save the weights of the pretrained model in a list
loaded_list = []
for item in encoder.variables:
  loaded_list.append((item.name,item.numpy()))<jupyter_output><empty_output><jupyter_text>Loading Weights in the pytorch modelChange now the model name with the one that matches the TF model loaded in the previous step.Different models have sligtly different behaviour.```pythonmodel_name = "modelA0" model_name = "modelA1"model_name = "modelA2"model_name = "modelA3"  model_name = "modelA4"   model_name = "modelA5" ```Change the loaded pytorch model ```pythonmodel = MoViNet(_C.MODEL.MoViNetA0, 600,causal = False, tf_like = True)model = MoViNet(_C.MODEL.MoViNetA1, 600,causal = False, tf_like = True)model = MoViNet(_C.MODEL.MoViNetA2, 600,causal = False, tf_like = True)model = MoViNet(_C.MODEL.MoViNetA3, 600,causal = False, tf_like = True)model = MoViNet(_C.MODEL.MoViNetA4, 600,causal = False, tf_like = True)model = MoViNet(_C.MODEL.MoViNetA5, 600,causal = False, tf_like = True)```<jupyter_code>model_name = "modelA0" 

def key_translate_base( k):
    k = (k
    .replace("classifier_head/head/conv3d/" ,"classifier.0.conv_1.conv3d.")
    .replace("classifier_head/classifier/conv3d/", "classifier.3.conv_1.conv3d.")
    .replace("se/se_reduce/conv3d/","se.fc1.conv_1.conv3d.")
    .replace("se/se_expand/conv3d/","se.fc2.conv_1.conv3d.")
    .replace("stem/stem/", "conv1.conv_1.")
    .replace("conv3d/", "conv3d.")
    .replace("kernel:0","weight")
    .replace("bias:0","bias")
    .replace("bn/gamma:0","norm.weight")
    .replace("bn/beta:0","norm.bias")
    .replace("bn/moving_mean:0","norm.running_mean")
    .replace("bn/moving_variance:0","norm.running_var")
    .replace("skip/skip_project/","res.1.conv_1.")
    .replace("expansion/","expand.conv_1.")
    .replace("feature/","deep.conv_1.")
    .replace("projection/","project.conv_1.")
    .replace("scale:0", "alpha")
    .replace("head/project/", "conv7.conv_1."))
    for i in range(5):
        for j in range(20):
            k=k.replace(f"b{i}/l{j}/bneck/", f"blocks.b{i}_l{j}.").replace(f"b{i}/l{j}/", f"blocks.b{i}_l{j}.")
    if (model_name == "modelA3" or model_name == "modelA5") and "b3_l0" in k:
        k = k.replace("res.1.","res.0.")
    return k

def key_translate_stream( k):
    k = (k.replace("feature/conv2d/depthwise_conv2d/depthwise_kernel:0","deep.conv_1.conv2d.weight")
         .replace("feature/conv2d_temporal/depthwise_conv2d_1/depthwise_kernel:0","deep.conv_2.conv2d.weight")
         
         .replace("feature/bn/","deep.conv_1.norm.")
         .replace("feature/bn_temporal/","deep.conv_2.norm.")
         .replace("expansion/conv2d/conv2d/","expand.conv_1.conv2d.")
         .replace("expansion/bn/","expand.conv_1.norm.")
         .replace("projection/conv2d/conv2d_3/","project.conv_1.conv2d.")
         .replace("projection/bn/","project.conv_1.norm.")
         .replace("se/se_reduce/conv2d/conv2d_1/","se.fc1.conv_1.conv2d.")
         .replace("se/se_expand/conv2d/conv2d_2/","se.fc2.conv_1.conv2d.")
         .replace("skip/skip_project/conv2d/conv2d_4/", "res.1.conv_1.conv2d.")
    .replace("skip/skip_project/bn/" ,"res.1.conv_1.norm.")
         .replace("classifier_head/head/conv2d/conv2d/" ,"classifier.0.conv_1.conv2d.")
    
         .replace("classifier_head/classifier/conv2d/conv2d_1/" ,"classifier.3.conv_1.conv2d.")
         .replace("head/project/bn/","conv7.conv_1.norm.")
         .replace("head/project/conv2d/conv2d/" ,"conv7.conv_1.conv2d.")
         .replace("stem/stem/bn/","conv1.conv_1.norm.")
         .replace("stem/stem/conv2d/conv2d/" ,"conv1.conv_1.conv2d.")
    .replace("kernel:0","weight")
    .replace("bias:0","bias")
    .replace("gamma:0","weight")
    .replace("beta:0","bias")
    .replace("moving_mean:0","running_mean")
    .replace("moving_variance:0","running_var")
    .replace("scale:0", "alpha")
        )
    for i in range(5):
        for j in range(20):
            k=k.replace(f"b{i}/l{j}/bneck/", f"blocks.b{i}_l{j}.").replace(f"b{i}/l{j}/", f"blocks.b{i}_l{j}.")
    if (model_name == "modelA3" or model_name == "modelA5") and "b3_l0" in k:
        k = k.replace("res.1.","res.0.")
    return k

def weight_translate( w):
    if len(w.shape)==5:
        w = rearrange(w, "d h w c_in c_out -> c_out c_in d h w")
    if len(w.shape)==4:
        #w = rearrange(w, "h w c_in c_out -> c_out c_in h w")
        if "feature" in name:
            w = rearrange(w, "h w c_out c_in-> c_out c_in h w")
        else:
            w = rearrange(w, "h w c_in c_out -> c_out c_in h w")
    return torch.tensor(w)

#creating the dictionary
param_dict = {key_translate_base(name ) : weight_translate( item) for i,(name,item) in enumerate(loaded_list)}

model = MoViNet(_C.MODEL.MoViNetA0, causal = False, num_classes = 600, tf_like = True)
#load the dictionary
model.load_state_dict(param_dict)
#save the model
torch.save(model.state_dict(), "./modelA0_statedict_v3")
from six.moves import urllib
from PIL import Image
image_url = 'https://upload.wikimedia.org/wikipedia/commons/8/84/Ski_Famille_-_Family_Ski_Holidays.jpg'
image_height = 172
image_width = 172
from io import BytesIO
import tensorflow as tf
import numpy as np

with urllib.request.urlopen(image_url) as f:
  image = Image.open(f).resize((image_height, image_width))
video = tf.reshape(np.array(image), [1, 1, image_height, image_width, 3])
video = tf.cast(video, tf.float32) / 255.
video_2 = rearrange(torch.from_numpy(video.numpy()), "b t h w c-> b c t h w")
model.eval()
model.clean_activation_buffers()
result = model(video_2)
torch.argmax(result)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.16837693308281543, 'fraction_numerical': 0.03396444572915161, 'mean_word_length': 3.996389891696751, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 14, 'https://': 8, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '2', 'dir': 'jupyter-structured-clean-dedup', 'id': '535736', 'license': 'MIT', 'n_tokens_mistral': 2698, 'n_tokens_neox': 2591, 'n_words': 447, 'path': 'weights/weight_load.ipynb', 'repo_name': 'patrick620/MoViNet-pytorch', 'size': '6919'}