From a26dcd9ef6f1d1c5684660d37d71f3898bfd6054 Mon Sep 17 00:00:00 2001
From: johnzylstra <johnmartinzylstra@gmail.com>
Date: Fri, 4 Mar 2022 02:18:24 -0500
Subject: [PATCH 1/8] Initial commit

Creating branch. Initial z head implementation with basic unit tests
---
 .../mesh_rcnn/modeling/heads/test_z_head.py   | 82 +++++++++++++++++++
 .../mesh_rcnn/modeling/heads/z_head.py        | 65 +++++++++++++++
 2 files changed, 147 insertions(+)
 create mode 100644 official/vision/beta/projects/mesh_rcnn/modeling/heads/test_z_head.py
 create mode 100644 official/vision/beta/projects/mesh_rcnn/modeling/heads/z_head.py

diff --git a/official/vision/beta/projects/mesh_rcnn/modeling/heads/test_z_head.py b/official/vision/beta/projects/mesh_rcnn/modeling/heads/test_z_head.py
new file mode 100644
index 00000000000..9a6ab4cdbe2
--- /dev/null
+++ b/official/vision/beta/projects/mesh_rcnn/modeling/heads/test_z_head.py
@@ -0,0 +1,82 @@
+import tensorflow as tf
+from absl.testing import parameterized
+from official.vision.beta.projects.mesh_rcnn.modeling.heads import z_head
+
+class ZHeadTest(parameterized.TestCase, tf.test.TestCase):
+        
+    def test_output_shape(self,
+        num_fc: int,
+        fc_dim: int,
+        cls_agnostic: bool,
+        num_classes: int,
+        **kwargs):
+
+        (batch_size, height, width, channels) = (64, 14, 14, 256)
+        
+        head = z_head.ZHead(num_fc, fc_dim, cls_agnostic, num_classes)
+        # For example use batch size 64, h,w of 14, and 256 channels
+        # NOTE: Don't actually know where the z-head attaches to architecture
+
+        input = tf.zeros((batch_size, height, width, channels))
+        output = head(input)
+        expected_output = tf.zeros((batch_size, num_classes))
+        self.assertAllEqual(output,expected_output)
+
+    def test_serialize_deserialize(self,
+        num_fc: int,
+        fc_dim: int,
+        cls_agnostic: bool,
+        num_classes: int,
+        **kwargs):
+
+        (batch_size, height, width, channels) = (64, 14, 14, 256)
+        head = z_head.ZHead(num_fc, fc_dim, cls_agnostic, num_classes)
+        
+        input = tf.zeros((batch_size, height, width, channels))
+        _ = head(input)
+
+        serialized = head.get_config()
+        deserialized = z_head.ZHead.from_config(serialized)
+
+        self.assertAllEqual(head.get_config(), deserialized.get_config())
+
+    def test_gradient_pass_through(self,
+        num_fc: int,
+        fc_dim: int,
+        cls_agnostic: bool,
+        num_classes: int,
+        **kwargs):
+
+        (batch_size, height, width, channels) = (64, 14, 14, 256)
+        head = z_head.ZHead(num_fc, fc_dim, cls_agnostic, num_classes)
+        
+        loss = tf.keras.losses.MeanSquaredError()
+        optimizer = tf.keras.optimizers.SGD()
+
+        init = tf.random_normal_initializer()
+        input_shape = (batch_size, height, width, channels)
+        x = tf.Variable(initial_value = init(shape=input_shape, dtype=tf.float32))
+
+        output_shape = head(x).shape
+        y = tf.Variable(initial_value = init(shape=output_shape, dtype=tf.float32))
+
+        with tf.GradientTape() as tape:
+            y_hat = head(x)
+            grad_loss = loss(y_hat, y)
+        grad = tape.gradient(grad_loss, head.trainable_variables)
+        optimizer.apply_gradients(zip(grad, head.trainable_variables))
+
+        self.assertNotIn(None, grad)
+
+    def test_build_from_config(self,
+        num_fc: int,
+        fc_dim: int,
+        cls_agnostic: bool,
+        num_classes: int,
+        **kwargs):
+        pass
+
+# zht = ZHeadTest()
+# zht.test_output_shape(2, 1024, False, 100)
+# zht.test_serialize_deserialize(2, 1024, False, 100)
+# zht.test_gradient_pass_through(2, 1024, False, 100)
diff --git a/official/vision/beta/projects/mesh_rcnn/modeling/heads/z_head.py b/official/vision/beta/projects/mesh_rcnn/modeling/heads/z_head.py
new file mode 100644
index 00000000000..785add9f770
--- /dev/null
+++ b/official/vision/beta/projects/mesh_rcnn/modeling/heads/z_head.py
@@ -0,0 +1,65 @@
+import tensorflow as tf
+
+class ZHead(tf.keras.layers.Layer):
+
+    def __init__(self,
+        num_fc: int,
+        fc_dim: int,
+        cls_agnostic: bool,
+        num_classes: int,
+        **kwargs):
+        """
+        Initialize Z-head
+        Args:
+            num_fc: number of fully connected layers
+            fc_dim: dimension of fully connected layers
+            cls_agnostic:
+            num_classes: number of prediction classes
+        """
+        super().__init__(**kwargs)
+
+        self._num_fc = num_fc
+        self._fc_dim = fc_dim
+        self._cls_agnostic = cls_agnostic
+        self._num_classes = num_classes
+
+    def build(self, input_shape: tf.TensorShape) -> None:
+        # INPUT SHAPE: N x H x W x C
+        self.flatten = tf.keras.layers.Flatten()
+
+        self.fcs = []
+        for k in range(self._num_fc):
+            fc = tf.keras.layers.Dense(self._fc_dim)
+            self.fcs.append(fc)
+        
+        num_z_reg_classes = 1 if self._cls_agnostic else self._num_classes
+        self.z_pred = tf.keras.layers.Dense(num_z_reg_classes, activation='relu')
+
+
+        # MAY HAVE TO DO WEIGHT INIT
+        # for layer in self.fcs:
+        #     weight_init.c2._xavier_fill(layer)
+
+        # nn.init.normal_(self.z_pred.weight, std=0.001)
+        # nn.init.constant_(self.z_pred.bias, 0)    
+
+    def call(self, x):
+        x = self.flatten(x)
+        for layer in self.fcs:
+            x = layer(x)
+        x = self.z_pred(x)
+        return x
+
+    def get_config(self):
+        """Get config dict of the ZHead layer."""
+        config = dict(
+            num_fc = self._num_fc,
+            fc_dim = self._fc_dim,
+            cls_agnostic = self._cls_agnostic,
+            num_classes = self._num_classes
+        )
+        return config
+
+    @classmethod
+    def from_config(cls, config):
+        return (cls(**config))

From 606c86e90a1d5505ff57729146430f11d55e629f Mon Sep 17 00:00:00 2001
From: johnzylstra <johnmartinzylstra@gmail.com>
Date: Fri, 4 Mar 2022 12:23:21 -0500
Subject: [PATCH 2/8] Some weight loading files

created some initial files for weightloading (based on those in mesh_rcnn_mesh_refinement_branch branch)
---
 .../utils/weight_utils/config_classes.py      | 111 +++++++++++++
 .../utils/weight_utils/config_data.py         |  35 ++++
 .../utils/weight_utils/load_weights.py        | 151 ++++++++++++++++++
 .../utils/weight_utils/test_load_weights.py   |  54 +++++++
 4 files changed, 351 insertions(+)
 create mode 100644 official/vision/beta/projects/mesh_rcnn/utils/weight_utils/config_classes.py
 create mode 100644 official/vision/beta/projects/mesh_rcnn/utils/weight_utils/config_data.py
 create mode 100644 official/vision/beta/projects/mesh_rcnn/utils/weight_utils/load_weights.py
 create mode 100644 official/vision/beta/projects/mesh_rcnn/utils/weight_utils/test_load_weights.py

diff --git a/official/vision/beta/projects/mesh_rcnn/utils/weight_utils/config_classes.py b/official/vision/beta/projects/mesh_rcnn/utils/weight_utils/config_classes.py
new file mode 100644
index 00000000000..bd8c5a2b361
--- /dev/null
+++ b/official/vision/beta/projects/mesh_rcnn/utils/weight_utils/config_classes.py
@@ -0,0 +1,111 @@
+"""Config Classes for weight-loading Mesh R-CNN"""
+
+from abc import ABC
+from dataclasses import dataclass, field
+from typing import Dict, List
+
+import numpy as np
+import tensorflow as tf
+
+
+class Config(ABC):
+  def get_weights(self):
+    """
+    This function generates the weights needed to be loaded into the layer.
+    """
+    return None
+
+  def load_weights(self, layer: tf.keras.layers.Layer) -> int:
+    """
+    Given a layer, this function retrives the weights for that layer in an
+    appropriate format, and loads them into the layer. Additionally,
+    the number of weights loaded are returned. If the weights are in an
+    incorrect format, a ValueError will be raised by set_weights().
+    """
+    weights = self.get_weights()
+    layer.set_weights(weights)
+    n_weights = 0
+
+    for w in weights:
+      n_weights += tf.size(w)
+    return n_weights
+
+@dataclass
+class ZHeadCFG(Config):
+  weights_dict: Dict = field(repr=False, default=None)
+  w0_weights: np.array = field(repr=False, default=None)
+  w0_bias: np.array = field(repr=False, default=None)
+  w1_weights: np.array = field(repr=False, default=None)
+  w1_bias: np.array = field(repr=False, default=None)
+
+  def __post_init__(self):
+    self.w0_weights = self.weights_dict['w0']['weight']
+    self.w0_bias = self.weights_dict['w0']['bias']
+    self.w1_weights = self.weights_dict['w1']['weight']
+    self.w1_bias = self.weights_dict['w1']['bias']
+
+    self.weights = [
+        self.w0_weights,
+        self.w0_bias,
+        self.w1_weights,
+        self.w1_bias
+    ]
+
+  def get_weights(self):
+    return self.weights
+
+# @dataclass
+# class meshRefinementStageCFG(Config):
+#   weights_dict: Dict = field(repr=False, default=None)
+#   bottleneck_weights: np.array = field(repr=False, default=None)
+#   bottleneck_bias: np.array = field(repr=False, default=None)
+#   verts_offset_weights: np.array = field(repr=False, default=None)
+#   verts_offset_bias: np.array = field(repr=False, default=None)
+#   gconvs_weights: np.array = field(repr=False, default=None)
+
+#   weights: np.array = field(repr=False, default=None)
+
+#   def __post_init__(self):
+#     bottleneck_weights = self.weights_dict['bottleneck']['weight']
+#     bottleneck_bias = self.weights_dict['bottleneck']['bias']
+#     verts_offset_weights = self.weights_dict['verts_offset']['weight']
+#     verts_offset_bias = self.weights_dict['verts_offset']['bias']
+
+#     self.weights = [
+#         bottleneck_weights,
+#         bottleneck_bias,
+#         verts_offset_weights,
+#         verts_offset_bias
+#     ]
+
+#     for stage in self.weights_dict['gconvs'].keys():
+#       gconvConfig = graphConvCFG(
+#           weights_dict=self.weights_dict['gconvs'][stage])
+#       self.weights += gconvConfig.get_weights()
+
+#   def get_weights(self):
+#     return self.weights
+
+# @dataclass
+# class graphConvCFG(Config):
+#   weights_dict: Dict = field(repr=False, default=None)
+#   w0_weights: np.array = field(repr=False, default=None)
+#   w0_bias: np.array = field(repr=False, default=None)
+#   w1_weights: np.array = field(repr=False, default=None)
+#   w1_bias: np.array = field(repr=False, default=None)
+
+#   def __post_init__(self):
+#     self.w0_weights = self.weights_dict['w0']['weight']
+#     self.w0_bias = self.weights_dict['w0']['bias']
+#     self.w1_weights = self.weights_dict['w1']['weight']
+#     self.w1_bias = self.weights_dict['w1']['bias']
+
+#     self.weights = [
+#         self.w0_weights,
+#         self.w0_bias,
+#         self.w1_weights,
+#         self.w1_bias
+#     ]
+
+#   def get_weights(self):
+#     return self.weights
\ No newline at end of file
diff --git a/official/vision/beta/projects/mesh_rcnn/utils/weight_utils/config_data.py b/official/vision/beta/projects/mesh_rcnn/utils/weight_utils/config_data.py
new file mode 100644
index 00000000000..b656aec251e
--- /dev/null
+++ b/official/vision/beta/projects/mesh_rcnn/utils/weight_utils/config_data.py
@@ -0,0 +1,35 @@
+"""Configs for model components."""
+
+from dataclasses import dataclass, field
+from typing import Dict
+
+# from official.vision.beta.projects.mesh_rcnn.utils.weight_utils.config_classes import \
+#     meshRefinementStageCFG
+from official.vision.beta.projects.mesh_rcnn.utils.weight_utils.config_classes import \
+    ZHeadCFG
+
+@dataclass
+class ZHeadConfigData():
+  weights_dict: Dict = field(repr=False, default=None)
+
+  def get_cfg_list(self, name):
+    if name == "pix3d":
+      return ZHeadCFG(weighs_dict=self.weights_dict)
+
+    else:
+      return []
+
+# @dataclass
+# class MeshHeadConfigData():
+#   weights_dict: Dict = field(repr=False, default=None)
+
+#   def get_cfg_list(self, name):
+#     if name == "pix3d":
+#       return [
+#           meshRefinementStageCFG(weights_dict=self.weights_dict['stages']['0']),
+#           meshRefinementStageCFG(weights_dict=self.weights_dict['stages']['1']),
+#           meshRefinementStageCFG(weights_dict=self.weights_dict['stages']['2']),
+#       ]
+
+#     else:
+#       return []
diff --git a/official/vision/beta/projects/mesh_rcnn/utils/weight_utils/load_weights.py b/official/vision/beta/projects/mesh_rcnn/utils/weight_utils/load_weights.py
new file mode 100644
index 00000000000..1bb4620fae5
--- /dev/null
+++ b/official/vision/beta/projects/mesh_rcnn/utils/weight_utils/load_weights.py
@@ -0,0 +1,151 @@
+"""
+This file contains functions used to load the Pytorch Mesh R-CNN checkpoint
+weights into the Tensorflow model.
+"""
+
+import numpy as np
+import tensorflow as tf
+from torch import load
+
+# from official.vision.beta.projects.mesh_rcnn.modeling.layers.nn_blocks import \
+#     MeshRefinementStage
+# from official.vision.beta.projects.mesh_rcnn.modeling.heads.z_head import \
+#     ZHead
+# from official.vision.beta.projects.mesh_rcnn.utils.weight_utils.config_data import \
+#     MeshHeadConfigData
+# from official.vision.beta.projects.mesh_rcnn.utils.weight_utils.config_data import \
+    # ZHeadConfigData
+
+
+def pth_to_dict(pth_path):
+  """ Converts a TF checkpoint into a nested dictionary of weights.
+  Args:
+  pth_path: String, indicating filepath of the Pytorch checkpoint
+  Returns:
+  Dictionary where the checkpoint weights are stored
+  Number of weights read
+  """
+  print("\nConverting model checkpoint from {} to weights dictionary\n".format(
+      pth_path))
+  
+  # Using cpu for now
+  pth_dict = load(pth_path, map_location='cpu')
+  weights_dict = {}
+  n_read = 0
+
+  for key in pth_dict["model"]:
+    keys = key.split('.')
+    curr_dict = weights_dict
+    for newKey in keys[:-1]:
+      if newKey not in curr_dict:
+        curr_dict[newKey] = {}
+      curr_dict = curr_dict[newKey]
+
+    np_tensor = pth_dict["model"][key].cpu().numpy()
+    # transpose from Pytorch to use in TF
+    if len(np_tensor.shape) == 4:
+      np_tensor = np.transpose(np_tensor, [2, 3, 1, 0])
+
+    if len(np_tensor.shape) == 2:
+      np_tensor = np.transpose(np_tensor, [1, 0])
+
+    tf_tensor = tf.convert_to_tensor(np_tensor)
+    curr_dict[keys[-1]] = tf_tensor
+
+    n_read += np.prod(np_tensor.shape)
+
+  print("Successfully read {} checkpoint weights\n".format(n_read))
+  return weights_dict, n_read
+
+
+# def get_mesh_head_layer_cfgs(weights_dict, mesh_head_name):
+#   """ Fetches the config classes for the mesh head.
+#   This function generates a list of config classes corresponding to
+#   each building block in the mesh head.
+#   Args:
+#     weights_dict: Dictionary that stores the backbone model weights.
+#     backbone_name: String, indicating the desired mesh head configuration.
+#   Returns:
+#     A list containing the config classes of the mesh head building block.
+#   """
+
+#   print("Fetching mesh head config classes for {}\n".format(mesh_head_name))
+#   cfgs = MeshHeadConfigData(weights_dict).get_cfg_list(mesh_head_name)
+#   return cfgs
+
+def get_zhead_layer_cfgs(weights_dict, zhead_name):
+  """ Fetches the config classes for the mesh head.
+  This function generates a list of config classes corresponding to
+  each building block in the mesh head.
+  Args:
+    weights_dict: Dictionary that stores the backbone model weights.
+    backbone_name: String, indicating the desired mesh head configuration.
+  Returns:
+    A list containing the config classes of the mesh head building block.
+  """
+
+  print("Fetching z head config classes for {}\n".format(zhead_name))
+  cfgs = ZHeadConfigData(weights_dict).get_cfg_list(zhead_name)
+  return cfgs
+
+# def load_weights_mesh_head(mesh_head, weights_dict, mesh_head_name):
+#   """ Loads the weights defined in the weights_dict into the backbone.
+#   This function loads the backbone weights by first fetching the necesary
+#   config classes for the backbone, then loads them in one by one for
+#   each layer that has weights associated with it.
+#   Args:
+#     backbone: keras.Model backbone.
+#     weights_dict: Dictionary that stores the backbone model weights.
+#     backbone_name: String, indicating the desired backbone configuration.
+#   Returns:
+#     Number of weights loaded in.
+#   """
+#   print("Loading mesh head weights\n")
+
+#   cfgs = get_mesh_head_layer_cfgs(weights_dict, mesh_head_name)
+#   n_weights_total = 0
+#   loaded_layers = 0
+
+#   i = 0
+#   for layer in mesh_head.layers:
+#     if isinstance(layer, (MeshRefinementStage)):
+#       n_weights = cfgs[i].load_weights(layer)
+#       print("Weights loaded for {}: {}".format(layer.name, n_weights))
+#       n_weights_total += n_weights
+#       loaded_layers += 1
+#       i += 1
+      
+#   print("{} Weights have been loaded for {} / {} layers\n".format(
+#       n_weights_total, loaded_layers, len(mesh_head.layers)))
+#   return n_weights_total
+
+def load_weights_zhead(mesh_head, weights_dict, zhead_name):
+  """ Loads the weights defined in the weights_dict into the backbone.
+  This function loads the backbone weights by first fetching the necesary
+  config classes for the backbone, then loads them in one by one for
+  each layer that has weights associated with it.
+  Args:
+    backbone: keras.Model backbone.
+    weights_dict: Dictionary that stores the backbone model weights.
+    backbone_name: String, indicating the desired backbone configuration.
+  Returns:
+    Number of weights loaded in.
+  """
+  print("Loading z head weights\n")
+
+  cfgs = get_zhead_layer_cfgs(weights_dict, zhead_name)
+  n_weights_total = 0
+  loaded_layers = 0
+
+  i = 0
+  for layer in mesh_head.layers:
+    if isinstance(layer, (ZHead)):
+      n_weights = cfgs[i].load_weights(layer)
+      print("Weights loaded for {}: {}".format(layer.name, n_weights))
+      n_weights_total += n_weights
+      loaded_layers += 1
+      i += 1
+      
+  print("{} Weights have been loaded for {} / {} layers\n".format(
+      n_weights_total, loaded_layers, len(mesh_head.layers)))
+  return n_weights_total
\ No newline at end of file
diff --git a/official/vision/beta/projects/mesh_rcnn/utils/weight_utils/test_load_weights.py b/official/vision/beta/projects/mesh_rcnn/utils/weight_utils/test_load_weights.py
new file mode 100644
index 00000000000..71ef0542323
--- /dev/null
+++ b/official/vision/beta/projects/mesh_rcnn/utils/weight_utils/test_load_weights.py
@@ -0,0 +1,54 @@
+from official.vision.beta.projects.mesh_rcnn.modeling.heads.z_head import \
+    ZHead
+# from ...modeling.heads.z_head import \
+#     ZHead
+# from official.vision.beta.projects.mesh_rcnn.ops.mesh_ops import \
+#     compute_mesh_shape
+# from official.vision.beta.projects.mesh_rcnn.utils.weight_utils.load_weights import (
+#     load_weights_zhead, pth_to_dict)
+
+from load_weights import load_weights_zhead, pth_to_dict
+
+PTH_PATH = r"C:\\Users\\johnm\\Downloads\\meshrcnn_R50.pth"
+
+def print_layer_names(layers_dict, offset=0):
+  if isinstance(layers_dict, dict):
+    for k in layers_dict.keys():
+      print(" " * offset + k)
+      print_layer_names(layers_dict[k], offset+2)
+
+def test_load_zhead():
+  weights_dict, n_read = pth_to_dict(PTH_PATH)
+  print(weights_dict.keys())
+  print(weights_dict['roi_heads'].keys())
+  print(weights_dict['roi_heads']['z_head'].keys())
+  print(weights_dict['roi_heads']['z_head']['z_pred'].keys())
+  print(weights_dict['roi_heads']['z_head']['z_pred']['weight'].shape)
+  # grid_dims = 24
+  # mesh_shapes = compute_mesh_shape(1, grid_dims)
+  # verts_shape, verts_mask_shape, faces_shape, faces_mask_shape = mesh_shapes
+  # backbone_shape = [1, 14, 14, 256]
+  # input_specs = {
+  #     'feature_map': backbone_shape,
+  #     'verts': verts_shape,
+  #     'verts_mask': verts_mask_shape,
+  #     'faces': faces_shape,
+  #     'faces_mask': faces_mask_shape
+  # }
+  # mesh_head = MeshHead(input_specs)
+  # mesh_head.summary()
+
+  input_specs = dict(
+    num_fc = 2,
+    fc_dim = 1024,
+    cls_agnostic = False,
+    num_classes = 9
+  )
+
+  zhead = ZHead.from_config(input_specs)
+
+  n_weights = load_weights_zhead(
+      zhead, weights_dict['roi_heads']['z_head'], 'pix3d')
+
+if __name__ == '__main__':
+  test_load_zhead()
\ No newline at end of file

From 8ce1db241d2e5a3c2ae7068e1dbd86b272195a6a Mon Sep 17 00:00:00 2001
From: johnzylstra <johnmartinzylstra@gmail.com>
Date: Tue, 29 Mar 2022 12:33:36 -0400
Subject: [PATCH 3/8] Diff Tested

Fully unit tested (test_z_head.py)
Weights loaded (test_load_weights.py)
Differential tested (code not provided)
---
 .../mesh_rcnn/modeling/heads/test_z_head.py   | 75 ++++++++++------
 .../mesh_rcnn/modeling/heads/z_head.py        | 57 +++++++-----
 .../utils/weight_utils/config_classes.py      | 90 +++++--------------
 .../utils/weight_utils/config_data.py         | 20 +----
 .../utils/weight_utils/load_weights.py        | 89 +++---------------
 .../utils/weight_utils/test_load_weights.py   | 21 +----
 6 files changed, 119 insertions(+), 233 deletions(-)

diff --git a/official/vision/beta/projects/mesh_rcnn/modeling/heads/test_z_head.py b/official/vision/beta/projects/mesh_rcnn/modeling/heads/test_z_head.py
index 9a6ab4cdbe2..8a82dbfbe4e 100644
--- a/official/vision/beta/projects/mesh_rcnn/modeling/heads/test_z_head.py
+++ b/official/vision/beta/projects/mesh_rcnn/modeling/heads/test_z_head.py
@@ -1,24 +1,39 @@
-import tensorflow as tf
+# Copyright 2021 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.
+"""Tests for Mesh R-CNN Heads."""
+
 from absl.testing import parameterized
-from official.vision.beta.projects.mesh_rcnn.modeling.heads import z_head
+import tensorflow as tf
+
+# from official.vision.beta.projects.mesh_rcnn.modeling.heads import z_head
+import z_head
 
 class ZHeadTest(parameterized.TestCase, tf.test.TestCase):
-        
+    '''Test for Mesh R-CNN Z head'''
+
     def test_output_shape(self,
         num_fc: int,
         fc_dim: int,
         cls_agnostic: bool,
-        num_classes: int,
-        **kwargs):
+        num_classes: int):
+        '''Check that Z head output is of correct shape'''
 
         (batch_size, height, width, channels) = (64, 14, 14, 256)
-        
         head = z_head.ZHead(num_fc, fc_dim, cls_agnostic, num_classes)
-        # For example use batch size 64, h,w of 14, and 256 channels
-        # NOTE: Don't actually know where the z-head attaches to architecture
 
-        input = tf.zeros((batch_size, height, width, channels))
-        output = head(input)
+        test_input = tf.zeros((batch_size, height, width, channels))
+        output = head(test_input)
         expected_output = tf.zeros((batch_size, num_classes))
         self.assertAllEqual(output,expected_output)
 
@@ -26,14 +41,13 @@ def test_serialize_deserialize(self,
         num_fc: int,
         fc_dim: int,
         cls_agnostic: bool,
-        num_classes: int,
-        **kwargs):
+        num_classes: int):
+        """Create a network object that sets all of its config options."""
 
         (batch_size, height, width, channels) = (64, 14, 14, 256)
         head = z_head.ZHead(num_fc, fc_dim, cls_agnostic, num_classes)
-        
-        input = tf.zeros((batch_size, height, width, channels))
-        _ = head(input)
+        test_input = tf.zeros((batch_size, height, width, channels))
+        _ = head(test_input)
 
         serialized = head.get_config()
         deserialized = z_head.ZHead.from_config(serialized)
@@ -44,25 +58,25 @@ def test_gradient_pass_through(self,
         num_fc: int,
         fc_dim: int,
         cls_agnostic: bool,
-        num_classes: int,
-        **kwargs):
+        num_classes: int):
+        '''Check that gradients are not None'''
 
         (batch_size, height, width, channels) = (64, 14, 14, 256)
         head = z_head.ZHead(num_fc, fc_dim, cls_agnostic, num_classes)
-        
+
         loss = tf.keras.losses.MeanSquaredError()
         optimizer = tf.keras.optimizers.SGD()
 
         init = tf.random_normal_initializer()
         input_shape = (batch_size, height, width, channels)
-        x = tf.Variable(initial_value = init(shape=input_shape, dtype=tf.float32))
+        test_input = tf.Variable(initial_value = init(shape=input_shape, dtype=tf.float32))
 
-        output_shape = head(x).shape
-        y = tf.Variable(initial_value = init(shape=output_shape, dtype=tf.float32))
+        output_shape = head(test_input).shape
+        test_output = tf.Variable(initial_value = init(shape=output_shape, dtype=tf.float32))
 
         with tf.GradientTape() as tape:
-            y_hat = head(x)
-            grad_loss = loss(y_hat, y)
+            y_hat = head(test_input)
+            grad_loss = loss(y_hat, test_output)
         grad = tape.gradient(grad_loss, head.trainable_variables)
         optimizer.apply_gradients(zip(grad, head.trainable_variables))
 
@@ -72,11 +86,14 @@ def test_build_from_config(self,
         num_fc: int,
         fc_dim: int,
         cls_agnostic: bool,
-        num_classes: int,
-        **kwargs):
+        num_classes: int):
+        '''build Z head from config'''
         pass
 
-# zht = ZHeadTest()
-# zht.test_output_shape(2, 1024, False, 100)
-# zht.test_serialize_deserialize(2, 1024, False, 100)
-# zht.test_gradient_pass_through(2, 1024, False, 100)
+if __name__ == '__main__':
+    print("Unit Testing Z Head")
+    zht = ZHeadTest()
+    zht.test_output_shape(2, 1024, False, 100)
+    zht.test_serialize_deserialize(2, 1024, False, 100)
+    zht.test_gradient_pass_through(2, 1024, False, 100)
+    print("Tests Successful")
diff --git a/official/vision/beta/projects/mesh_rcnn/modeling/heads/z_head.py b/official/vision/beta/projects/mesh_rcnn/modeling/heads/z_head.py
index 785add9f770..06486922613 100644
--- a/official/vision/beta/projects/mesh_rcnn/modeling/heads/z_head.py
+++ b/official/vision/beta/projects/mesh_rcnn/modeling/heads/z_head.py
@@ -1,7 +1,22 @@
+# Copyright 2021 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.
+"""Mesh R-CNN Heads."""
+
 import tensorflow as tf
 
 class ZHead(tf.keras.layers.Layer):
-
+    '''Depth prediction Z Head for Mesh R-CNN model'''
     def __init__(self,
         num_fc: int,
         fc_dim: int,
@@ -16,7 +31,7 @@ def __init__(self,
             cls_agnostic:
             num_classes: number of prediction classes
         """
-        super().__init__(**kwargs)
+        super(ZHead, self).__init__(**kwargs)
 
         self._num_fc = num_fc
         self._fc_dim = fc_dim
@@ -24,31 +39,28 @@ def __init__(self,
         self._num_classes = num_classes
 
     def build(self, input_shape: tf.TensorShape) -> None:
-        # INPUT SHAPE: N x H x W x C
+        '''Build Z Head'''
         self.flatten = tf.keras.layers.Flatten()
 
         self.fcs = []
-        for k in range(self._num_fc):
-            fc = tf.keras.layers.Dense(self._fc_dim)
-            self.fcs.append(fc)
-        
+        for _ in range(self._num_fc):
+            layer = tf.keras.layers.Dense(self._fc_dim,
+                activation='relu',
+                kernel_initializer='he_uniform')
+            self.fcs.append(layer)
         num_z_reg_classes = 1 if self._cls_agnostic else self._num_classes
-        self.z_pred = tf.keras.layers.Dense(num_z_reg_classes, activation='relu')
-
-
-        # MAY HAVE TO DO WEIGHT INIT
-        # for layer in self.fcs:
-        #     weight_init.c2._xavier_fill(layer)
-
-        # nn.init.normal_(self.z_pred.weight, std=0.001)
-        # nn.init.constant_(self.z_pred.bias, 0)    
+        pred_init = tf.keras.initializers.RandomNormal(stddev=0.001)
+        self.z_pred = tf.keras.layers.Dense(num_z_reg_classes,
+            kernel_initializer=pred_init,
+            bias_initializer='zeros')
 
-    def call(self, x):
-        x = self.flatten(x)
+    def call(self, features):
+        '''Forward pass of Z head'''
+        out = self.flatten(features)
         for layer in self.fcs:
-            x = layer(x)
-        x = self.z_pred(x)
-        return x
+            out = layer(out)
+        out = self.z_pred(out)
+        return out
 
     def get_config(self):
         """Get config dict of the ZHead layer."""
@@ -62,4 +74,5 @@ def get_config(self):
 
     @classmethod
     def from_config(cls, config):
-        return (cls(**config))
+        '''Initialize Z head from config'''
+        return cls(**config)
diff --git a/official/vision/beta/projects/mesh_rcnn/utils/weight_utils/config_classes.py b/official/vision/beta/projects/mesh_rcnn/utils/weight_utils/config_classes.py
index bd8c5a2b361..2cb7096bf61 100644
--- a/official/vision/beta/projects/mesh_rcnn/utils/weight_utils/config_classes.py
+++ b/official/vision/beta/projects/mesh_rcnn/utils/weight_utils/config_classes.py
@@ -33,79 +33,31 @@ def load_weights(self, layer: tf.keras.layers.Layer) -> int:
 @dataclass
 class ZHeadCFG(Config):
   weights_dict: Dict = field(repr=False, default=None)
-  w0_weights: np.array = field(repr=False, default=None)
-  w0_bias: np.array = field(repr=False, default=None)
-  w1_weights: np.array = field(repr=False, default=None)
-  w1_bias: np.array = field(repr=False, default=None)
+  weights: List = field(repr=False, default=None)
+
+  fc1_weights: np.array = field(repr=False, default=None)
+  fc1_bias: np.array = field(repr=False, default=None)
+  fc2_weights: np.array = field(repr=False, default=None)
+  fc2_bias: np.array = field(repr=False, default=None)
+  pred_weights: np.array = field(repr=False, default=None)
+  pred_bias: np.array = field(repr=False, default=None)
 
   def __post_init__(self):
-    self.w0_weights = self.weights_dict['w0']['weight']
-    self.w0_bias = self.weights_dict['w0']['bias']
-    self.w1_weights = self.weights_dict['w1']['weight']
-    self.w1_bias = self.weights_dict['w1']['bias']
+    self.fc1_weights = self.weights_dict['z_fc1']['weight']
+    self.fc1_bias = self.weights_dict['z_fc1']['bias']
+    self.fc2_weights = self.weights_dict['z_fc2']['weight']
+    self.fc2_bias = self.weights_dict['z_fc2']['bias']
+    self.pred_weights = self.weights_dict['z_pred']['weight']
+    self.pred_bias = self.weights_dict['z_pred']['bias']
 
     self.weights = [
-        self.w0_weights,
-        self.w0_bias,
-        self.w1_weights,
-        self.w1_bias
+        self.fc1_weights,
+        self.fc1_bias,
+        self.fc2_weights,
+        self.fc2_bias,
+        self.pred_weights,
+        self.pred_bias
     ]
 
   def get_weights(self):
-    return self.weights
-
-# @dataclass
-# class meshRefinementStageCFG(Config):
-#   weights_dict: Dict = field(repr=False, default=None)
-#   bottleneck_weights: np.array = field(repr=False, default=None)
-#   bottleneck_bias: np.array = field(repr=False, default=None)
-#   verts_offset_weights: np.array = field(repr=False, default=None)
-#   verts_offset_bias: np.array = field(repr=False, default=None)
-#   gconvs_weights: np.array = field(repr=False, default=None)
-
-#   weights: np.array = field(repr=False, default=None)
-
-#   def __post_init__(self):
-#     bottleneck_weights = self.weights_dict['bottleneck']['weight']
-#     bottleneck_bias = self.weights_dict['bottleneck']['bias']
-#     verts_offset_weights = self.weights_dict['verts_offset']['weight']
-#     verts_offset_bias = self.weights_dict['verts_offset']['bias']
-
-#     self.weights = [
-#         bottleneck_weights,
-#         bottleneck_bias,
-#         verts_offset_weights,
-#         verts_offset_bias
-#     ]
-
-#     for stage in self.weights_dict['gconvs'].keys():
-#       gconvConfig = graphConvCFG(
-#           weights_dict=self.weights_dict['gconvs'][stage])
-#       self.weights += gconvConfig.get_weights()
-
-#   def get_weights(self):
-#     return self.weights
-
-# @dataclass
-# class graphConvCFG(Config):
-#   weights_dict: Dict = field(repr=False, default=None)
-#   w0_weights: np.array = field(repr=False, default=None)
-#   w0_bias: np.array = field(repr=False, default=None)
-#   w1_weights: np.array = field(repr=False, default=None)
-#   w1_bias: np.array = field(repr=False, default=None)
-
-#   def __post_init__(self):
-#     self.w0_weights = self.weights_dict['w0']['weight']
-#     self.w0_bias = self.weights_dict['w0']['bias']
-#     self.w1_weights = self.weights_dict['w1']['weight']
-#     self.w1_bias = self.weights_dict['w1']['bias']
-
-#     self.weights = [
-#         self.w0_weights,
-#         self.w0_bias,
-#         self.w1_weights,
-#         self.w1_bias
-#     ]
-
-#   def get_weights(self):
-#     return self.weights
\ No newline at end of file
+    return self.weights
\ No newline at end of file
diff --git a/official/vision/beta/projects/mesh_rcnn/utils/weight_utils/config_data.py b/official/vision/beta/projects/mesh_rcnn/utils/weight_utils/config_data.py
index b656aec251e..dea0fd70e63 100644
--- a/official/vision/beta/projects/mesh_rcnn/utils/weight_utils/config_data.py
+++ b/official/vision/beta/projects/mesh_rcnn/utils/weight_utils/config_data.py
@@ -3,10 +3,9 @@
 from dataclasses import dataclass, field
 from typing import Dict
 
-# from official.vision.beta.projects.mesh_rcnn.utils.weight_utils.config_classes import \
-#     meshRefinementStageCFG
 from official.vision.beta.projects.mesh_rcnn.utils.weight_utils.config_classes import \
     ZHeadCFG
+# from config_classes import ZHeadCFG
 
 @dataclass
 class ZHeadConfigData():
@@ -14,22 +13,7 @@ class ZHeadConfigData():
 
   def get_cfg_list(self, name):
     if name == "pix3d":
-      return ZHeadCFG(weighs_dict=self.weights_dict)
+      return ZHeadCFG(weights_dict=self.weights_dict)
 
     else:
       return []
-
-# @dataclass
-# class MeshHeadConfigData():
-#   weights_dict: Dict = field(repr=False, default=None)
-
-#   def get_cfg_list(self, name):
-#     if name == "pix3d":
-#       return [
-#           meshRefinementStageCFG(weights_dict=self.weights_dict['stages']['0']),
-#           meshRefinementStageCFG(weights_dict=self.weights_dict['stages']['1']),
-#           meshRefinementStageCFG(weights_dict=self.weights_dict['stages']['2']),
-#       ]
-
-#     else:
-#       return []
diff --git a/official/vision/beta/projects/mesh_rcnn/utils/weight_utils/load_weights.py b/official/vision/beta/projects/mesh_rcnn/utils/weight_utils/load_weights.py
index 1bb4620fae5..5920c02c648 100644
--- a/official/vision/beta/projects/mesh_rcnn/utils/weight_utils/load_weights.py
+++ b/official/vision/beta/projects/mesh_rcnn/utils/weight_utils/load_weights.py
@@ -6,15 +6,12 @@
 import numpy as np
 import tensorflow as tf
 from torch import load
-
-# from official.vision.beta.projects.mesh_rcnn.modeling.layers.nn_blocks import \
-#     MeshRefinementStage
-# from official.vision.beta.projects.mesh_rcnn.modeling.heads.z_head import \
-#     ZHead
-# from official.vision.beta.projects.mesh_rcnn.utils.weight_utils.config_data import \
-#     MeshHeadConfigData
-# from official.vision.beta.projects.mesh_rcnn.utils.weight_utils.config_data import \
-    # ZHeadConfigData
+from official.vision.beta.projects.mesh_rcnn.modeling.heads.z_head import \
+    ZHead
+# from z_head import ZHead
+from official.vision.beta.projects.mesh_rcnn.utils.weight_utils.config_data import \
+    ZHeadConfigData
+# from config_data import ZHeadConfigData
 
 
 def pth_to_dict(pth_path):
@@ -28,8 +25,7 @@ def pth_to_dict(pth_path):
   print("\nConverting model checkpoint from {} to weights dictionary\n".format(
       pth_path))
   
-  # Using cpu for now
-  pth_dict = load(pth_path, map_location='cpu')
+  pth_dict = load(pth_path, map_location='cpu') # may need to delete cpu arg
   weights_dict = {}
   n_read = 0
 
@@ -58,21 +54,6 @@ def pth_to_dict(pth_path):
   return weights_dict, n_read
 
 
-# def get_mesh_head_layer_cfgs(weights_dict, mesh_head_name):
-#   """ Fetches the config classes for the mesh head.
-#   This function generates a list of config classes corresponding to
-#   each building block in the mesh head.
-#   Args:
-#     weights_dict: Dictionary that stores the backbone model weights.
-#     backbone_name: String, indicating the desired mesh head configuration.
-#   Returns:
-#     A list containing the config classes of the mesh head building block.
-#   """
-
-#   print("Fetching mesh head config classes for {}\n".format(mesh_head_name))
-#   cfgs = MeshHeadConfigData(weights_dict).get_cfg_list(mesh_head_name)
-#   return cfgs
-
 def get_zhead_layer_cfgs(weights_dict, zhead_name):
   """ Fetches the config classes for the mesh head.
   This function generates a list of config classes corresponding to
@@ -88,64 +69,20 @@ def get_zhead_layer_cfgs(weights_dict, zhead_name):
   cfgs = ZHeadConfigData(weights_dict).get_cfg_list(zhead_name)
   return cfgs
 
-# def load_weights_mesh_head(mesh_head, weights_dict, mesh_head_name):
-#   """ Loads the weights defined in the weights_dict into the backbone.
-#   This function loads the backbone weights by first fetching the necesary
-#   config classes for the backbone, then loads them in one by one for
-#   each layer that has weights associated with it.
-#   Args:
-#     backbone: keras.Model backbone.
-#     weights_dict: Dictionary that stores the backbone model weights.
-#     backbone_name: String, indicating the desired backbone configuration.
-#   Returns:
-#     Number of weights loaded in.
-#   """
-#   print("Loading mesh head weights\n")
-
-#   cfgs = get_mesh_head_layer_cfgs(weights_dict, mesh_head_name)
-#   n_weights_total = 0
-#   loaded_layers = 0
 
-#   i = 0
-#   for layer in mesh_head.layers:
-#     if isinstance(layer, (MeshRefinementStage)):
-#       n_weights = cfgs[i].load_weights(layer)
-#       print("Weights loaded for {}: {}".format(layer.name, n_weights))
-#       n_weights_total += n_weights
-#       loaded_layers += 1
-#       i += 1
-      
-#   print("{} Weights have been loaded for {} / {} layers\n".format(
-#       n_weights_total, loaded_layers, len(mesh_head.layers)))
-#   return n_weights_total
-
-def load_weights_zhead(mesh_head, weights_dict, zhead_name):
+def load_weights_zhead(zhead, weights_dict, zhead_name):
   """ Loads the weights defined in the weights_dict into the backbone.
   This function loads the backbone weights by first fetching the necesary
   config classes for the backbone, then loads them in one by one for
   each layer that has weights associated with it.
   Args:
-    backbone: keras.Model backbone.
-    weights_dict: Dictionary that stores the backbone model weights.
-    backbone_name: String, indicating the desired backbone configuration.
+    zhead: keras.Model
+    weights_dict: Dictionary that stores the zhead model weights.
+    zhead_name: String, indicating the desired zhead configuration.
   Returns:
     Number of weights loaded in.
   """
   print("Loading z head weights\n")
-
   cfgs = get_zhead_layer_cfgs(weights_dict, zhead_name)
-  n_weights_total = 0
-  loaded_layers = 0
-
-  i = 0
-  for layer in mesh_head.layers:
-    if isinstance(layer, (ZHead)):
-      n_weights = cfgs[i].load_weights(layer)
-      print("Weights loaded for {}: {}".format(layer.name, n_weights))
-      n_weights_total += n_weights
-      loaded_layers += 1
-      i += 1
-      
-  print("{} Weights have been loaded for {} / {} layers\n".format(
-      n_weights_total, loaded_layers, len(mesh_head.layers)))
-  return n_weights_total
\ No newline at end of file
+  cfgs.load_weights(zhead)
+  return
\ No newline at end of file
diff --git a/official/vision/beta/projects/mesh_rcnn/utils/weight_utils/test_load_weights.py b/official/vision/beta/projects/mesh_rcnn/utils/weight_utils/test_load_weights.py
index 71ef0542323..9b46b90f170 100644
--- a/official/vision/beta/projects/mesh_rcnn/utils/weight_utils/test_load_weights.py
+++ b/official/vision/beta/projects/mesh_rcnn/utils/weight_utils/test_load_weights.py
@@ -1,11 +1,7 @@
 from official.vision.beta.projects.mesh_rcnn.modeling.heads.z_head import \
     ZHead
-# from ...modeling.heads.z_head import \
-#     ZHead
-# from official.vision.beta.projects.mesh_rcnn.ops.mesh_ops import \
-#     compute_mesh_shape
-# from official.vision.beta.projects.mesh_rcnn.utils.weight_utils.load_weights import (
-#     load_weights_zhead, pth_to_dict)
+from official.vision.beta.projects.mesh_rcnn.utils.weight_utils.load_weights import (
+    load_weights_zhead, pth_to_dict)
 
 from load_weights import load_weights_zhead, pth_to_dict
 
@@ -24,19 +20,6 @@ def test_load_zhead():
   print(weights_dict['roi_heads']['z_head'].keys())
   print(weights_dict['roi_heads']['z_head']['z_pred'].keys())
   print(weights_dict['roi_heads']['z_head']['z_pred']['weight'].shape)
-  # grid_dims = 24
-  # mesh_shapes = compute_mesh_shape(1, grid_dims)
-  # verts_shape, verts_mask_shape, faces_shape, faces_mask_shape = mesh_shapes
-  # backbone_shape = [1, 14, 14, 256]
-  # input_specs = {
-  #     'feature_map': backbone_shape,
-  #     'verts': verts_shape,
-  #     'verts_mask': verts_mask_shape,
-  #     'faces': faces_shape,
-  #     'faces_mask': faces_mask_shape
-  # }
-  # mesh_head = MeshHead(input_specs)
-  # mesh_head.summary()
 
   input_specs = dict(
     num_fc = 2,

From 30d1fd3d42386f6a539d3f39ba72bd83623c82d4 Mon Sep 17 00:00:00 2001
From: johnzylstra <johnmartinzylstra@gmail.com>
Date: Tue, 29 Mar 2022 15:24:43 -0400
Subject: [PATCH 4/8] Config and minor fixes

added config in mesh_rcnn.py
Fixed fc layer weight initializer
---
 .../projects/mesh_rcnn/configs/mesh_rcnn.py   | 27 +++++++++++++++++++
 .../mesh_rcnn/modeling/heads/test_z_head.py   |  8 ------
 .../mesh_rcnn/modeling/heads/z_head.py        |  5 +++-
 .../utils/weight_utils/test_load_weights.py   |  2 ++
 4 files changed, 33 insertions(+), 9 deletions(-)
 create mode 100644 official/vision/beta/projects/mesh_rcnn/configs/mesh_rcnn.py

diff --git a/official/vision/beta/projects/mesh_rcnn/configs/mesh_rcnn.py b/official/vision/beta/projects/mesh_rcnn/configs/mesh_rcnn.py
new file mode 100644
index 00000000000..823beec555b
--- /dev/null
+++ b/official/vision/beta/projects/mesh_rcnn/configs/mesh_rcnn.py
@@ -0,0 +1,27 @@
+# Copyright 2021 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.
+"""Mesh R-CNN configuration definition."""
+
+import dataclasses
+
+from official.modeling import hyperparams  # type: ignore
+
+
+@dataclasses.dataclass
+class ZHead(hyperparams.Config):
+    """Parameterization for the Mesh R-CNN Z Head."""
+    num_fc: int = 2
+    fc_dim: int = 1024
+    cls_agnostic: bool = False
+    num_classes: int = 9
\ No newline at end of file
diff --git a/official/vision/beta/projects/mesh_rcnn/modeling/heads/test_z_head.py b/official/vision/beta/projects/mesh_rcnn/modeling/heads/test_z_head.py
index 8a82dbfbe4e..b5dfa63d49e 100644
--- a/official/vision/beta/projects/mesh_rcnn/modeling/heads/test_z_head.py
+++ b/official/vision/beta/projects/mesh_rcnn/modeling/heads/test_z_head.py
@@ -82,14 +82,6 @@ def test_gradient_pass_through(self,
 
         self.assertNotIn(None, grad)
 
-    def test_build_from_config(self,
-        num_fc: int,
-        fc_dim: int,
-        cls_agnostic: bool,
-        num_classes: int):
-        '''build Z head from config'''
-        pass
-
 if __name__ == '__main__':
     print("Unit Testing Z Head")
     zht = ZHeadTest()
diff --git a/official/vision/beta/projects/mesh_rcnn/modeling/heads/z_head.py b/official/vision/beta/projects/mesh_rcnn/modeling/heads/z_head.py
index 06486922613..7f2495a4b4f 100644
--- a/official/vision/beta/projects/mesh_rcnn/modeling/heads/z_head.py
+++ b/official/vision/beta/projects/mesh_rcnn/modeling/heads/z_head.py
@@ -44,9 +44,12 @@ def build(self, input_shape: tf.TensorShape) -> None:
 
         self.fcs = []
         for _ in range(self._num_fc):
+            fc_init = tf.keras.initializers.VarianceScaling(scale=1.,
+              mode='fan_in',
+              distribution='truncated_normal')
             layer = tf.keras.layers.Dense(self._fc_dim,
                 activation='relu',
-                kernel_initializer='he_uniform')
+                kernel_initializer=fc_init)
             self.fcs.append(layer)
         num_z_reg_classes = 1 if self._cls_agnostic else self._num_classes
         pred_init = tf.keras.initializers.RandomNormal(stddev=0.001)
diff --git a/official/vision/beta/projects/mesh_rcnn/utils/weight_utils/test_load_weights.py b/official/vision/beta/projects/mesh_rcnn/utils/weight_utils/test_load_weights.py
index 9b46b90f170..576087f60a6 100644
--- a/official/vision/beta/projects/mesh_rcnn/utils/weight_utils/test_load_weights.py
+++ b/official/vision/beta/projects/mesh_rcnn/utils/weight_utils/test_load_weights.py
@@ -1,3 +1,4 @@
+'''Test weightloading for Z head'''
 from official.vision.beta.projects.mesh_rcnn.modeling.heads.z_head import \
     ZHead
 from official.vision.beta.projects.mesh_rcnn.utils.weight_utils.load_weights import (
@@ -8,6 +9,7 @@
 PTH_PATH = r"C:\\Users\\johnm\\Downloads\\meshrcnn_R50.pth"
 
 def print_layer_names(layers_dict, offset=0):
+  '''print names of layers from dictionary'''
   if isinstance(layers_dict, dict):
     for k in layers_dict.keys():
       print(" " * offset + k)

From 5b7d0ec2a5889866822368a740ac106b2a7d55e9 Mon Sep 17 00:00:00 2001
From: johnzylstra <johnmartinzylstra@gmail.com>
Date: Tue, 29 Mar 2022 15:29:06 -0400
Subject: [PATCH 5/8] Cosmetic

newlines at the end of files
---
 .../projects/mesh_rcnn/utils/weight_utils/config_classes.py     | 2 +-
 .../beta/projects/mesh_rcnn/utils/weight_utils/load_weights.py  | 2 +-
 .../projects/mesh_rcnn/utils/weight_utils/test_load_weights.py  | 2 +-
 3 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/official/vision/beta/projects/mesh_rcnn/utils/weight_utils/config_classes.py b/official/vision/beta/projects/mesh_rcnn/utils/weight_utils/config_classes.py
index 2cb7096bf61..38fc16da192 100644
--- a/official/vision/beta/projects/mesh_rcnn/utils/weight_utils/config_classes.py
+++ b/official/vision/beta/projects/mesh_rcnn/utils/weight_utils/config_classes.py
@@ -60,4 +60,4 @@ def __post_init__(self):
     ]
 
   def get_weights(self):
-    return self.weights
\ No newline at end of file
+    return self.weights
diff --git a/official/vision/beta/projects/mesh_rcnn/utils/weight_utils/load_weights.py b/official/vision/beta/projects/mesh_rcnn/utils/weight_utils/load_weights.py
index 5920c02c648..9a9efd92ada 100644
--- a/official/vision/beta/projects/mesh_rcnn/utils/weight_utils/load_weights.py
+++ b/official/vision/beta/projects/mesh_rcnn/utils/weight_utils/load_weights.py
@@ -85,4 +85,4 @@ def load_weights_zhead(zhead, weights_dict, zhead_name):
   print("Loading z head weights\n")
   cfgs = get_zhead_layer_cfgs(weights_dict, zhead_name)
   cfgs.load_weights(zhead)
-  return
\ No newline at end of file
+  return
diff --git a/official/vision/beta/projects/mesh_rcnn/utils/weight_utils/test_load_weights.py b/official/vision/beta/projects/mesh_rcnn/utils/weight_utils/test_load_weights.py
index 576087f60a6..cb236e13139 100644
--- a/official/vision/beta/projects/mesh_rcnn/utils/weight_utils/test_load_weights.py
+++ b/official/vision/beta/projects/mesh_rcnn/utils/weight_utils/test_load_weights.py
@@ -36,4 +36,4 @@ def test_load_zhead():
       zhead, weights_dict['roi_heads']['z_head'], 'pix3d')
 
 if __name__ == '__main__':
-  test_load_zhead()
\ No newline at end of file
+  test_load_zhead()

From f5a1e26f1fd58e7dcc5bb8d72dedc11706b86eb3 Mon Sep 17 00:00:00 2001
From: johnzylstra <49248070+johnzylstra@users.noreply.github.com>
Date: Sun, 3 Apr 2022 20:10:05 -0400
Subject: [PATCH 6/8] Aligning Function Arguments

Aligning function arguments to be in same column
---
 .../projects/mesh_rcnn/modeling/heads/z_head.py    | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/official/vision/beta/projects/mesh_rcnn/modeling/heads/z_head.py b/official/vision/beta/projects/mesh_rcnn/modeling/heads/z_head.py
index 7f2495a4b4f..e257a082114 100644
--- a/official/vision/beta/projects/mesh_rcnn/modeling/heads/z_head.py
+++ b/official/vision/beta/projects/mesh_rcnn/modeling/heads/z_head.py
@@ -18,11 +18,11 @@
 class ZHead(tf.keras.layers.Layer):
     '''Depth prediction Z Head for Mesh R-CNN model'''
     def __init__(self,
-        num_fc: int,
-        fc_dim: int,
-        cls_agnostic: bool,
-        num_classes: int,
-        **kwargs):
+                 num_fc: int,
+                 fc_dim: int,
+                 cls_agnostic: bool,
+                 num_classes: int,
+                 **kwargs):
         """
         Initialize Z-head
         Args:
@@ -45,8 +45,8 @@ def build(self, input_shape: tf.TensorShape) -> None:
         self.fcs = []
         for _ in range(self._num_fc):
             fc_init = tf.keras.initializers.VarianceScaling(scale=1.,
-              mode='fan_in',
-              distribution='truncated_normal')
+                mode='fan_in',
+                distribution='truncated_normal')
             layer = tf.keras.layers.Dense(self._fc_dim,
                 activation='relu',
                 kernel_initializer=fc_init)

From a2df15e0b72f20f586e73a7bb81a3266465edcd2 Mon Sep 17 00:00:00 2001
From: johnzylstra <49248070+johnzylstra@users.noreply.github.com>
Date: Sun, 3 Apr 2022 20:12:07 -0400
Subject: [PATCH 7/8] test_z_head main

test_z_head main function now says "tf.test.main()"
---
 .../beta/projects/mesh_rcnn/modeling/heads/test_z_head.py  | 7 +------
 1 file changed, 1 insertion(+), 6 deletions(-)

diff --git a/official/vision/beta/projects/mesh_rcnn/modeling/heads/test_z_head.py b/official/vision/beta/projects/mesh_rcnn/modeling/heads/test_z_head.py
index b5dfa63d49e..266ca8d133f 100644
--- a/official/vision/beta/projects/mesh_rcnn/modeling/heads/test_z_head.py
+++ b/official/vision/beta/projects/mesh_rcnn/modeling/heads/test_z_head.py
@@ -83,9 +83,4 @@ def test_gradient_pass_through(self,
         self.assertNotIn(None, grad)
 
 if __name__ == '__main__':
-    print("Unit Testing Z Head")
-    zht = ZHeadTest()
-    zht.test_output_shape(2, 1024, False, 100)
-    zht.test_serialize_deserialize(2, 1024, False, 100)
-    zht.test_gradient_pass_through(2, 1024, False, 100)
-    print("Tests Successful")
+    tf.test.main()

From 158fbbaf73ffca67cb52e950f6b8694d7e9624bc Mon Sep 17 00:00:00 2001
From: johnzylstra <johnmartinzylstra@gmail.com>
Date: Mon, 4 Apr 2022 16:41:20 -0400
Subject: [PATCH 8/8] Z head test params

Add parameters to z-head test
---
 .../projects/mesh_rcnn/modeling/heads/test_z_head.py     | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/official/vision/beta/projects/mesh_rcnn/modeling/heads/test_z_head.py b/official/vision/beta/projects/mesh_rcnn/modeling/heads/test_z_head.py
index 266ca8d133f..b6769fb1aa9 100644
--- a/official/vision/beta/projects/mesh_rcnn/modeling/heads/test_z_head.py
+++ b/official/vision/beta/projects/mesh_rcnn/modeling/heads/test_z_head.py
@@ -16,9 +16,16 @@
 from absl.testing import parameterized
 import tensorflow as tf
 
-# from official.vision.beta.projects.mesh_rcnn.modeling.heads import z_head
 import z_head
 
+@parameterized.named_parameters(
+    {'testcase_name': 'pix3d_params',
+     'num_fc': 2,
+     'fc_dim': 1024,
+     'cls_agnostic': False,
+     'num_classes': 9
+    }
+)
 class ZHeadTest(parameterized.TestCase, tf.test.TestCase):
     '''Test for Mesh R-CNN Z head'''