Tensorflow migration problem regardless defining the operations tensorflow.python.framework.ops

Question

I am running a model from github and I already encountered several errors with pathing etc. After fixing this, I think the main error for now is tensorflow. This repo was probably done when TF 1.x, and now with the change to TF 2, I might need to migrate everything.

Mainly, I get the following error:

    @ops.RegisterShape('ApproxMatch')
AttributeError: module 'tensorflow.python.framework.ops' has no attribute 'RegisterShape'

in:

import tensorflow as tf
from tensorflow.python.framework import ops
import os.path as osp

base_dir = osp.dirname(osp.abspath(__file__))

approxmatch_module = tf.load_op_library(osp.join(base_dir, 'tf_approxmatch_so.so'))


def approx_match(xyz1,xyz2):
    '''
input:
    xyz1 : batch_size * #dataset_points * 3
    xyz2 : batch_size * #query_points * 3
returns:
    match : batch_size * #query_points * #dataset_points
    '''
    return approxmatch_module.approx_match(xyz1,xyz2)
ops.NoGradient('ApproxMatch')
#@tf.RegisterShape('ApproxMatch')
@ops.RegisterShape('ApproxMatch')
def _approx_match_shape(op):
    shape1=op.inputs[0].get_shape().with_rank(3)
    shape2=op.inputs[1].get_shape().with_rank(3)
    return [tf.TensorShape([shape1.dims[0],shape2.dims[1],shape1.dims[1]])]

2 main things I am not understanding:

1. I read that this probably will make me have to create the ops C++ routines but at the same time, I can see that these are done in here: @ops.RegisterShape('ApproxMatch') . Like this and this with REGISTER_OP(...).SetShapeFn(...). But I dont think I am understanding the process and have seen other questions with the same, but no real implementation/answer.
2. If I go to the location of the tf_approxmatch shared library ( approxmatch_module = tf.load_op_library(osp.join(base_dir, 'tf_approxmatch_so.so')) ), I cannot open it or edit it with gedit, so I am assuming I am not supposed to change anything in there (?).

There are py, cpp and cu files in that folder (I already did make yesterday and everything ran smoothly).

__init__.py     tf_approxmatch.cu.o   tf_nndistance.cu.o
makefile        tf_approxmatch.py     tf_nndistance.py
__pycache__     tf_approxmatch_so.so  tf_nndistance_so.so
tf_approxmatch.cpp  tf_nndistance.cpp
tf_approxmatch.cu   tf_nndistance.cu

My main guess is that I should register the operation of RegisterShape somehow in the cpp file as it already has some registered ops, but I am a bit lost because I am not even sure if I am understanding the problem I have. I will just show the first lines of the file:

#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_kernel.h"
#include <algorithm>
#include <vector>
#include <math.h>
using namespace tensorflow;
REGISTER_OP("ApproxMatch")
    .Input("xyz1: float32")
    .Input("xyz2: float32")
    .Output("match: float32");
REGISTER_OP("MatchCost")
    .Input("xyz1: float32")
    .Input("xyz2: float32")
    .Input("match: float32")
    .Output("cost: float32");
REGISTER_OP("MatchCostGrad")
    .Input("xyz1: float32")
    .Input("xyz2: float32")
    .Input("match: float32")
    .Output("grad1: float32")
    .Output("grad2: float32");

Answer 1

Disclaimer: Unless no other choice, I would strongly recommend sticking to TensorFlow 1.x. Migrating code from TF 1.x to 2.x can be incredibly time consuming.

---

Registering a shape is done in c++ using SetShapeFn and not in python since TF 1.0. However, the private python API was kept in TF 1.x (I presume for backward compatibility reasons), but was completely removed in TF 2.0.

The Create an Op guide is really useful to migrate the code in that case, and I would really recommend reading it.

First of all, why is it needed to register a shape? It's needed for shape inference, a feature that allow TensorFlow to know the shape of the inputs and outputs in the computation graph without running actual code. Shape Inference allows, for example, error handling when trying to use an op on Tensors that don't have a compatible shape.

In your specific case, you need to convert the python code that uses ops.RegisterShape to c++ code using SetShapeFn. Thankfully, the github repository you are working with provides comments that are going to be helpful.

Let's start with the approx_match function. The python code is the following:

def approx_match(xyz1,xyz2):
    '''
input:
    xyz1 : batch_size * #dataset_points * 3
    xyz2 : batch_size * #query_points * 3
returns:
    match : batch_size * #query_points * #dataset_points
    '''
    return approxmatch_module.approx_match(xyz1,xyz2)
@ops.RegisterShape('ApproxMatch')
def _approx_match_shape(op):
    shape1=op.inputs[0].get_shape().with_rank(3)
    shape2=op.inputs[1].get_shape().with_rank(3)
    return [tf.TensorShape([shape1.dims[0],shape2.dims[1],shape1.dims[1]])]

Reading the code and the comments, we understand the following:

• There is 2 inputs xyz1, and xyz2
• xyz1 has the shape (batch_size, dataset_points, 3)
• xyz2 has the shape (batch_size, query_points, 3)
• There is 1 output: match
• match has the shape (batch_size, query_points, dataset_points)

This would translate to the following C++ code:

#include "tensorflow/core/framework/shape_inference.h"
using namespace tensorflow;
REGISTER_OP("ApproxMatch")
    .Input("xyz1: float32")
    .Input("xyz2: float32")
    .Output("match: float32")
    .SetShapeFn([](shape_inference::InferenceContext* c) {
        shape_inference::ShapeHandle xyz1_shape = c->input(0);
        shape_inference::ShapeHandle xyz2_shape = c->input(1);
        // batch_size is the first dimension 
        shape_inference::DimensionHandle batch_size = c->Dim(xyz1_shape, 0);
        // dataset_points points is the 2nd dimension of the first input
        shape_inference::DimensionHandle dataset_points = c->Dim(xyz1_shape, 1);
        // query_points points is the 2nd dimension of the second input
        shape_inference::DimensionHandle query_points = c->Dim(xyz2_shape, 1);
        // Creating the new shape (batch_size, query_points, dataset_points)
        // and setting it to the output
        c->set_output(0, c->MakeShape({batch_size, query_points, dataset_points})); 
        // Returning a status telling that everything went well
        return Status::OK();    
    });

Warning: This code does not contain any error handling (for example, checking that the first dimensions of the 2 inputs are the same, or that the last dimension of both inputs is 3). I let that as an exercise to the reader, you could look at the aforementioned guide or directly to the source code of some ops to have an idea on how to do error handling, with, for example, the macro TF_RETURN_IF_ERROR.

---

The same steps can be applied to the match_cost function, which could look like that:

REGISTER_OP("MatchCost")
    .Input("xyz1: float32")
    .Input("xyz2: float32")
    .Input("match: float32")
    .Output("cost: float32")
    .SetShapeFn([](shape_inference::InferenceContext* c) {
        shape_inference::DimensionHandle batch_size = c->Dim(c->input(0), 0);
        c->set_output(0, c->Vector(batch_size));
        return Status::OK();        
    });

---

You need then to recompile the so library with the makefile included in the project. You might have to change some flags, for example, TF 2.8 uses the c++ standard c++14 instead of c++11, so the flag -std=c++14 is needed. Once the library is compiled, you can test importing it in python:

>>> import tensorflow as tf
>>> approxmatch_module = tf.load_op_library('./tf_approxmatch_so.so')
>>> a = tf.random.uniform((10,20,3))
>>> b = tf.random.uniform((10,50,3))
>>> c = approxmatch_module.approx_match(a,b)
>>> c.shape
TensorShape([10, 50, 20])

Answer 2

according to tensorflow's release log, RegisterShape is deprecated, and you should use SetShapeFn to define the shape when registering your operator in c++ source file.