From aaab34a72c13ae3974fcd80664d0494fe2adfcbb Mon Sep 17 00:00:00 2001 From: Jan-Frederik Schulte Date: Tue, 22 Oct 2024 13:21:25 -0400 Subject: [PATCH] fix softmax parsing in pytorch and add test --- hls4ml/converters/pytorch/core.py | 19 ++++++++++++++----- test/pytest/test_pytorch_api.py | 9 +++++++++ 2 files changed, 23 insertions(+), 5 deletions(-) diff --git a/hls4ml/converters/pytorch/core.py b/hls4ml/converters/pytorch/core.py index c56857715a..d6b46d93a3 100644 --- a/hls4ml/converters/pytorch/core.py +++ b/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 diff --git a/test/pytest/test_pytorch_api.py b/test/pytest/test_pytorch_api.py index fee7b9a3aa..6d558d4b35 100644 --- a/test/pytest/test_pytorch_api.py +++ b/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(),