saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-09 17:28:58 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Continue memory planning when unknown shape tensor is encountered. (#6413)

Browse source
This commit is contained in:
M. Zeeshan Siddiqui 2021-01-22 14:37:53 -08:00 committed by GitHub
parent 61ecf52c24
commit 3c3d36334b
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 11 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

16
onnxruntime/core/framework/session_state.cc
View file

@ -389,6 +389,7 @@ void TryCalculateSizeFromResolvedShape(int ml_value_idx, std::unordered_map<int,
} // namespace
// If this function fails NO memory planning will take place, hence lets ONLY FAIL and stop training where warranted, example SIZE overflow.
Status SessionState::GeneratePatternGroupCache(const std::vector<std::reference_wrapper<const TensorShape>>& input_shape,
const std::vector<int>& feed_mlvalue_idxs,
MemoryPatternGroup* output,
@ -425,12 +426,17 @@ Status SessionState::GeneratePatternGroupCache(const std::vector<std::reference_
auto* arg = node->OutputDefs()[i];
size_t is_resolved = 0;
std::vector<int64_t> resolved_shape;
ORT_RETURN_IF_ERROR(TryResolveShape(arg, map, is_resolved, resolved_shape));
// Store all valid resolved shapes. They will be queried in, for example,
// Recv operator to bypass the dependency of output shapes on inputs.
if (is_resolved != 0) {
resolved_shapes[ml_value_idx] = resolved_shape;
// Tensors whose shape cannot be resolved statically will be allocated at runtime.
if (TryResolveShape(arg, map, is_resolved, resolved_shape).IsOK()) {
// Store all valid resolved shapes. They will be queried in, for example,
// Recv operator to bypass the dependency of output shapes on inputs.
if (is_resolved != 0) {
resolved_shapes[ml_value_idx] = resolved_shape;
}
} else {
LOGS(logger_, INFO) << "[Static memory planning] Could not resolve shape for tensor with ML index "
<< ml_value_idx << ", will allocate dynamically.";
}
}
}