support temperature for softmax

python/mxnet/contrib/onnx/mx2onnx/_op_translations.py

@@ -859,15 +859,21 @@ def convert_softmax(node, **kwargs):
     name, input_nodes, attrs = get_inputs(node, kwargs)
 
     axis = int(attrs.get("axis", -1))
-    temperature = attrs.get("temperature", None)
-    if temperature and float(temperature) != 1.0:
-        raise NotImplementedError("Temperature is not supported for now.")
-    use_length = attrs.get("use_length", None)
+    temperature = str(attrs.get("temperature", 'None'))
+    if temperature == 'None':
+        temperature = 1.
+    else:
+        temperature = float(temperature)
+
+    use_length = str(attrs.get("use_length", 'None'))
     input_type = kwargs["in_type"]
+    dtype = onnx.mapping.TENSOR_TYPE_TO_NP_TYPE[input_type]
     data = input_nodes[0]
 
     nodes = [
-        make_node("Exp", [data], [name+"_exp_out"]),
+        create_tensor([temperature], name+"_tmp", kwargs["initializer"], dtype=dtype),
+        make_node("Div", [data, name+"_tmp"], [name+'_data']),
+        make_node("Exp", [name+'_data'], [name+"_exp_out"]),
         make_node("ReduceSum", [name+"_exp_out"], [name+"_rsum_out"], axes=[axis], keepdims=1)
     ]
     if len(input_nodes) == 1:

tests/python-pytest/onnx/test_operators.py

@@ -338,17 +338,18 @@ def test_onnx_export_cast(tmp_path, src_dtype, dst_dtype, shape):
 
 
 @pytest.mark.parametrize('dtype', ['float16', 'float32'])
-def test_onnx_export_softmax(tmp_path, dtype):
+@pytest.mark.parametrize('temperature', [.1, 1., 10.])
+def test_onnx_export_softmax(tmp_path, dtype, temperature):
     x = mx.nd.random.uniform(0, 1, (2, 3, 4), dtype=dtype)
-    M1 = def_model('softmax')
+    M1 = def_model('softmax', temperature=temperature)
     op_export_test('softmax_1', M1, [x], tmp_path)
-    M2 = def_model('softmax', use_length=True, axis=0)
+    M2 = def_model('softmax', use_length=True, axis=0, temperature=temperature)
     l2 = mx.nd.array([[2,0,2,1],[1,1,2,1], [0,0,0,1]], dtype=int)
     op_export_test('softmax_2', M2, [x, l2], tmp_path)
-    M3 = def_model('softmax', use_length=True, axis=-1)
+    M3 = def_model('softmax', use_length=True, axis=-1, temperature=temperature)
     l3 = mx.nd.array([[2,0,4],[0,0,0]], dtype=int)
     op_export_test('softmax_3', M3, [x, l3], tmp_path)
-    M4 = def_model('softmax', use_length=True, axis=1)
+    M4 = def_model('softmax', use_length=True, axis=1, temperature=temperature)
     l4 = mx.nd.array([[2,0,3,1],[0,1,0,0]], dtype=int)
     op_export_test('softmax_4', M4, [x, l4], tmp_path)