CODEWITHSUNDEEP
pythontensorflow
Swas_99
Swas_99

Reputation: 2390

TensorFlow: How to define shape of output in SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c)

The output tensor is defined as follow:

// Grab the input tensor
const Tensor& I_tensor = context->input(0);
const Tensor& W_tensor = context->input(1);
auto Input = I_tensor.flat<float>();
auto Weight = W_tensor.flat<float>();
// OP_REQUIRES(context, iA_tensor.dims()==2 && iB_tensor.dims()==2);

int B = I_tensor.dim_size(1);
int nH= I_tensor.dim_size(2);
int nW= I_tensor.dim_size(3);
int C = I_tensor.dim_size(4);
int K = W_tensor.dim_size(2);

// Create an output tensor
Tensor* O_tensor = NULL;
OP_REQUIRES_OK(context, context->allocate_output(0, TensorShape{B, 2*nH, 2*nW, K}, &O_tensor));

I tried using the same code inside the SetShapeFn, but it doesn't work.

.SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
  // // const Tensor& I_tensor = c->input(0); //This doesn't work
  // // const Tensor& W_tensor = c->input(1); //This doesn't work

  ::tensorflow::shape_inference::ShapeHandle I = c->input(0);
  ::tensorflow::shape_inference::ShapeHandle W = c->input(1);
  //I have no idea how to get the desired shape from I and W

  return Status::OK();
})

I have tried to find more info about ShapeHandle and InferenceContext, in order to get the desired result but have failed to do so. I will be grateful if someone can help me with this.

Thanks!

Upvotes: 2

Views: 969

Answers (1)

Vojtech Molek
Vojtech Molek

Reputation: 108

This is what I have used to get dimensions info and to create new ones:

auto N = c->Dim(c->input(0), 0);
auto H = c->Dim(c->input(0), 1);
auto W = c->Dim(c->input(0), 2);
auto C = c->Dim(c->input(0), 3);

// save stride so we can use it to calculate output shape
std::vector<float> stride;
c->GetAttr("stride", &stride);
H = c->MakeDim(round(c->Value(H)/stride[0]));
W = c->MakeDim(round(c->Value(W)/stride[1]));
c->set_output(0, c->MakeShape({N, H, W, C}));

Upvotes: 1

Related Questions