mardi 22 mai 2018

SetShapeFn change output shape of 4 dimensions

REGISTER_OP("Forward")
.Attr("F: {float, int32}")
.Input("img0: F")
.SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {

  // c->set_output(0, c->Matrix((3,4,3,3)));
  c->set_output(0, c->input(0)); // N C H W
  return Status::OK();})
);

How to change the output channels C? I can only set matrix(x,y) of two dimensions, but I need 4 dimensions as the c->input(0) tensor.

Aucun commentaire:

Enregistrer un commentaire