models/official/projects/movinet/tools/convert_3d_2plus1d.py

# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Converts '3d_2plus1d' checkpoints into '2plus1d'."""

from absl import app
from absl import flags
import tensorflow as tf, tf_keras

from official.projects.movinet.modeling import movinet
from official.projects.movinet.modeling import movinet_model

flags.DEFINE_string(
    'input_checkpoint_path', None,
    'Checkpoint path to load.')
flags.DEFINE_string(
    'output_checkpoint_path', None,
    'Export path to save the saved_model file.')
flags.DEFINE_string(
    'model_id', 'a0', 'MoViNet model name.')
flags.DEFINE_string(
    'se_type', '2plus3d', 'MoViNet model SE type.')
flags.DEFINE_bool(
    'causal', True, 'Run the model in causal mode.')
flags.DEFINE_bool(
    'use_positional_encoding', False,
    'Whether to use positional encoding (only applied when causal=True).')
flags.DEFINE_integer(
    'num_classes', 600, 'The number of classes for prediction.')
flags.DEFINE_bool(
    'verify_output', False, 'Verify the output matches between the models.')

FLAGS = flags.FLAGS


def main(_) -> None:
  backbone_2plus1d = movinet.Movinet(
      model_id=FLAGS.model_id,
      causal=FLAGS.causal,
      conv_type='2plus1d',
      se_type=FLAGS.se_type,
      use_positional_encoding=FLAGS.use_positional_encoding)
  model_2plus1d = movinet_model.MovinetClassifier(
      backbone=backbone_2plus1d,
      num_classes=FLAGS.num_classes)
  model_2plus1d.build([1, 1, 1, 1, 3])

  backbone_3d_2plus1d = movinet.Movinet(
      model_id=FLAGS.model_id,
      causal=FLAGS.causal,
      conv_type='3d_2plus1d',
      se_type=FLAGS.se_type,
      use_positional_encoding=FLAGS.use_positional_encoding)
  model_3d_2plus1d = movinet_model.MovinetClassifier(
      backbone=backbone_3d_2plus1d,
      num_classes=FLAGS.num_classes)
  model_3d_2plus1d.build([1, 1, 1, 1, 3])

  checkpoint = tf.train.Checkpoint(model=model_3d_2plus1d)
  status = checkpoint.restore(FLAGS.input_checkpoint_path)
  status.assert_existing_objects_matched()

  # Ensure both models have the same weights
  weights = []
  for var_2plus1d, var_3d_2plus1d in zip(
      model_2plus1d.get_weights(), model_3d_2plus1d.get_weights()):
    if var_2plus1d.shape == var_3d_2plus1d.shape:
      weights.append(var_3d_2plus1d)
    else:
      if var_3d_2plus1d.shape[0] == 1:
        weight = var_3d_2plus1d[0]
      else:
        weight = var_3d_2plus1d[:, 0]
      if weight.shape[-1] != var_2plus1d.shape[-1]:
        # Transpose any depthwise kernels (conv3d --> depthwise_conv2d)
        weight = tf.transpose(weight, perm=(0, 1, 3, 2))
      weights.append(weight)
  model_2plus1d.set_weights(weights)

  if FLAGS.verify_output:
    inputs = tf.random.uniform([1, 6, 64, 64, 3], dtype=tf.float32)

    logits_2plus1d = model_2plus1d(inputs)
    logits_3d_2plus1d = model_3d_2plus1d(inputs)

    if tf.reduce_mean(logits_2plus1d - logits_3d_2plus1d) > 1e-5:
      raise ValueError('Bad conversion, model outputs do not match.')

  save_checkpoint = tf.train.Checkpoint(
      model=model_2plus1d, backbone=backbone_2plus1d)
  save_checkpoint.save(FLAGS.output_checkpoint_path)


if __name__ == '__main__':
  flags.mark_flag_as_required('input_checkpoint_path')
  flags.mark_flag_as_required('output_checkpoint_path')
  app.run(main)