fix softmax parsing in pytorch and add test

Author: JanFSchulte

hls4ml/converters/pytorch/core.py

@@ -61,17 +61,21 @@ def parse_activation_layer(operation, layer_name, input_names, input_shapes, nod
             layer['class_name'] = 'ThresholdedReLU'
             layer['activation'] = 'ThresholdedReLU'
             if layer['activ_param'] < 0:
-                raise Exception('negative threshold values not supported')
-
-        if hasattr(node, 'dim'):
+                raise Exception('negative threshold values not supported') 
+        if hasattr(class_object, 'dim'):
             layer['axis'] = class_object.dim
+            if layer['class_name'] == 'Softmax' and layer['axis'] is None:
+                 layer['axis'] = -1
+            if 'IOType' in config:     
+                if  layer['class_name'] == 'Softmax' and config['IOType'] == 'io_stream' and layer['axis'] != -1:
+                    raise Exception('dim needs to be -1 for io_stream')  
     else:
         if layer['class_name'] in ['ReLU', 'Sigmoid', 'Tanh']:
             layer['class_name'] = 'Activation'
         if layer['class_name'] == 'LeakyReLU':
             layer['activ_param'] = node.kwargs['negative_slope']
         if layer['class_name'] == 'ELU':
-            layer['activ_param'] = node.kwargs['alpha']
+            layer['activ_param'] = node.kwargs['alpha']      
         if layer['class_name'] == 'Threshold':
             layer['activ_param'] = node.args[1]
             if layer['activ_param'] < 0:
@@ -80,7 +84,12 @@ def parse_activation_layer(operation, layer_name, input_names, input_shapes, nod
             layer['activation'] = 'ThresholdedReLU'
         if 'dim' in node.kwargs:
             layer['axis'] = node.kwargs['dim']
-
+            if layer['class_name'] == 'Softmax' and layer['axis'] is None:
+                 layer['axis'] = -1
+            if 'IOType' in config:     
+                if  layer['class_name'] == 'Softmax' and config['IOType'] == 'io_stream' and layer['axis'] != -1:
+                    raise Exception('dim needs to be -1 for io_stream')  
+                             
     output_shape = input_shapes[0]
     return layer, output_shape
 

test/pytest/test_pytorch_api.py

@@ -63,6 +63,7 @@ def test_linear(backend, io_type):
 @pytest.mark.parametrize(
     "activation_function",
     [
+        nn.Softmax(dim=-1),
         nn.ReLU(),
         nn.Tanh(),
         nn.LeakyReLU(negative_slope=1.0),
@@ -74,6 +75,7 @@ def test_linear(backend, io_type):
 )
 @pytest.mark.parametrize('backend', ['Vivado', 'Vitis', 'Quartus'])
 @pytest.mark.parametrize('io_type', ['io_parallel', 'io_stream'])
+
 def test_activations(activation_function, backend, io_type):
     model = torch.nn.Sequential(nn.Linear(1, 1), activation_function).to()
     model.eval()
@@ -118,6 +120,12 @@ def __init__(self):
     def forward(self, x):
         return nn.functional.relu(x)
 
+class SoftmaxModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        return nn.functional.softmax(x,dim=-1)
 
 class TanHModel(nn.Module):
     def __init__(self):
@@ -162,6 +170,7 @@ def forward(self, x):
 @pytest.mark.parametrize(
     "activation_function",
     [
+        SoftmaxModel(),
         ReLuModel(),
         TanHModel(),
         LeakyReLuModel(),