tfkit::tf Namespace Reference

Functions
bool	HasAttr (const tensorflow::NodeDef &node, const std::string &attr_name)
tensorflow::DataType	GetDataTypeAttr (const tensorflow::NodeDef &node, const std::string &attr_name)
tensorflow::TensorProto *	GetTensorAttr (tensorflow::NodeDef &node, const std::string &attr_name)
int	GetElementCount (const tensorflow::TensorShapeProto &)
	GetElementCount returns -1 for rank-0 tensor shape.

Function Documentation

GetDataTypeAttr()

tensorflow::DataType tfkit::tf::GetDataTypeAttr	(	const tensorflow::NodeDef &	node,
		const std::string &	attr_name
	)

Definition at line 58 of file Support.cpp.

{
  assert(HasAttr(node, attr_name));
  const auto &attr = node.attr().at(attr_name);
  assert(attr.value_case() == tensorflow::AttrValue::kType);
  return attr.type();
}

References HasAttr().

GetElementCount()

int tfkit::tf::GetElementCount

(

const tensorflow::TensorShapeProto &

shape

)

GetElementCount returns -1 for rank-0 tensor shape.

Definition at line 74 of file Support.cpp.

{
  int count = -1;
 
  for (auto &d : shape.dim())
  {
    if (d.size() == 0)
    {
      count = 0;
      break;
    }
    if (count == -1)
      count = 1;
 
    count *= d.size();
  }
  return count;
}

GetTensorAttr()

tensorflow::TensorProto * tfkit::tf::GetTensorAttr	(	tensorflow::NodeDef &	node,
		const std::string &	attr_name
	)

Definition at line 66 of file Support.cpp.

{
  assert(HasAttr(node, attr_name));
  tensorflow::AttrValue &attr = node.mutable_attr()->at(attr_name);
  assert(attr.value_case() == tensorflow::AttrValue::kTensor);
  return attr.mutable_tensor();
}

References HasAttr().

HasAttr()

bool tfkit::tf::HasAttr	(	const tensorflow::NodeDef &	node,
		const std::string &	attr_name
	)

Definition at line 53 of file Support.cpp.

{
  return node.attr().count(attr_name) > 0;
}

Referenced by GetDataTypeAttr(), and GetTensorAttr().