Skip to content

[AMP] Register some new ops #9849

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

Merged
merged 3 commits into from
Jan 7, 2022
Merged

[AMP] Register some new ops #9849

Changes from 2 commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
d35b7aa
add new mixed precision ops
AndrewZhaoLuo Jan 6, 2022
ae0a5f8
remove space
AndrewZhaoLuo Jan 6, 2022
7a87ede
Update python/tvm/relay/transform/mixed_precision.py
AndrewZhaoLuo Jan 6, 2022
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
9 changes: 8 additions & 1 deletion python/tvm/relay/transform/mixed_precision.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,7 @@
"scatter",
"full",
"dyn.full",
"nn.depth_to_space",
# Comparison
"less",
"greater",
Expand All @@ -88,13 +89,19 @@
"min",
"maximum",
"minimum",
"argmax",
"argmin",
"nn.relu",
"nn.leaky_relu",
"nn.prelu",
"nn.dropout",
# Complicated activations which saturate in a narrow range
"sigmoid",
"tanh",
"fast_tanh", # Some coefficients outside of representable range, but probably ok
"fast_exp",
"fast_erf",
"clip", # Usually safe, my result in oddity if clip greater than fp16 range
# Pooling operations
"nn.max_pool1d",
"nn.max_pool2d",
Expand All @@ -108,6 +115,7 @@
"nn.adaptive_max_pool1d",
"nn.adaptive_max_pool2d",
"nn.adaptive_max_pool3d",
"image.resize2d",
]
DEFAULT_NEVER_LIST = [
# In general if |f(x)| >> |x| for expected inputs then put the op here.
Expand All @@ -134,7 +142,6 @@
"nn.layer_norm",
]


# Returns a decorator which registers for every given op, the function under FTVMMixedPrecisionConversionType
def register_func_to_op_list(list_ops: List):
def decorator(func):
Expand Down