Skip to content

Avoid requiring static shapes for Size and Split ops #757

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
adk9 wants to merge 5 commits into
base: master
Choose a base branch
from
Open

Avoid requiring static shapes for Size and Split ops #757

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
2405e06
Fix Size translation to support dynamic inputs
Dec 16, 2020
08bd3fb
Avoid fixing negative dims in Split since it's handled by IE
Dec 16, 2020
ef52dda
Compute shape statically if input is static
Dec 17, 2020
fbcab2a
Handle negative axis for split
Dec 17, 2020
ab8035a
Merge branch 'master' into adk9/fix-static-translations
adk9 Dec 17, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 18 additions & 16 deletions ngraph_bridge/ngraph_builder.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2122,21 +2122,27 @@ static Status TranslateSizeOp(const Node* op, const std::vector<const Tensor*>&,

DataType dtype;
TF_RETURN_IF_ERROR(GetNodeAttr(op->attrs(), "out_type", &dtype));

// Size has an attribute to specify output, int32 or int64
ng::element::Type type;
TF_RETURN_IF_ERROR(util::TFDataTypeToNGraphElementType(dtype, &type));

auto ng_input_shape = ng_input.get_shape();
int64 result = 1;
for (auto dim : ng_input_shape) {
result *= dim;
ngraph::Output<ngraph::Node> ng_result;
if (ng_input.get_partial_shape().is_static()) {
auto ng_input_shape = ng_input.get_shape();
int64 result = 1;
for (auto dim : ng_input_shape) {
result *= dim;
}
// make a scalar with value equals to result
ng_result = ConstructNgNode<opset::Constant>(op->name(), type, ng::Shape(0),
std::vector<int64>({result}));
} else {
auto ng_input_shape =
ConstructNgNode<opset::ShapeOf>(op->name(), ng_input, type);
auto ng_axis = ConstructNgNode<opset::Constant>(
op->name(), ngraph::element::i64, ngraph::Shape{1}, vector<int64_t>{0});
ng_result =
ConstructNgNode<opset::ReduceProd>(op->name(), ng_input_shape, ng_axis);
}

// make a scalar with value equals to result
auto ng_result = ConstructNgNode<opset::Constant>(
op->name(), type, ng::Shape(0), std::vector<int64>({result}));

SaveNgOp(ng_op_map, op->name(), ng_result);
return Status::OK();
}
Expand Down Expand Up @@ -2258,15 +2264,11 @@ static Status TranslateSplitOp(
int32 num_split;
TF_RETURN_IF_ERROR(GetNodeAttr(op->attrs(), "num_split", &num_split));

ng::Shape shape = ng_input.get_shape();
int rank = shape.size();

std::vector<int> split_dim_vec;
TF_RETURN_IF_ERROR(
GetStaticInputVector(op, 0, static_input_map, &split_dim_vec));
int split_dim = split_dim_vec[0] + (split_dim_vec[0] < 0 ? (int64)rank : 0);
auto ng_split_dim = ConstructNgNode<opset::Constant>(
op->name(), ng::element::u64, ng::Shape{}, split_dim);
op->name(), ng::element::i32, ng::Shape{}, split_dim_vec[0]);
auto ng_split = make_shared<opset::Split>(ng_input, ng_split_dim, num_split);

for (int i = 0; i < num_split; ++i) {
Expand Down
4 changes: 3 additions & 1 deletion ngraph_bridge/ngraph_mark_for_clustering.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -736,7 +736,9 @@ GetTFToNgOpMap() {
{"Sigmoid", {std::make_shared<opset::Sigmoid>()}},
{"Sin", {std::make_shared<opset::Sin>()}},
{"Sinh", {std::make_shared<opset::Sinh>()}},
{"Size", {constant}},
{"Size",
{std::make_shared<opset::ShapeOf>(),
std::make_shared<opset::ReduceProd>()}},
{"Sign", {std::make_shared<opset::Sign>()}},
{"Slice", {constant, std::make_shared<opset::StridedSlice>()}},
{"Snapshot", {}},
Expand Down