Skip to content

Improved ONNX export for relu6 layer #2986

Description

@bas-aarts

the current ONNX export for relu6 generates a min/max pair (

@torch_op("aten::relu6", trace_only=True)
def aten_relu6(self: TReal) -> TReal:
"""relu6(Tensor self) -> Tensor"""
six = op.CastLike(op.Constant(value_int=6), self)
return op.Min(op.Relu(self), six)
)

given that ONNX has a Clip layer, Why not use

    zero = op.CastLike(op.Constant(value_int=0), self)
    six = op.CastLike(op.Constant(value_int=6), self)
    return op.Clip(self, zero, six)

instead?

Sure, the sequence is easy to optimize downstream, but generating the Clip layer here makes a lot of sense.

Metadata

Metadata

Labels

No labels
No labels

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions