fastmachinelearning / hls4ml

Machine learning on FPGAs using HLS
https://fastmachinelearning.org/hls4ml
Apache License 2.0
1.24k stars 402 forks source link

Bug fix for named nn.Sequential in pytorch parser #848

Closed JanFSchulte closed 1 year ago

JanFSchulte commented 1 year ago

Parsing of nn.Sequentials that are named members of a model class results in a naming convention for the tensors in the state_dict of the model different from what the parser expects, since it was so far tested only on unnamed nn.Sequentials. This PR catches this and adjusts the name of the tensors we are importing from the state_dict accordingly. A test is added to ensure that we keep parsing both cases successfully.

Type of change

For a new feature or function, please create an issue first to discuss it with us before submitting a pull request.

Note: Please delete options that are not relevant.

Tests

To reproduce, this will fail with this PR:


import torch.nn as nn

from hls4ml.converters import convert_from_pytorch_model
from hls4ml.utils.config import config_from_pytorch_model

#simple model with namend sequential
class SeqModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.layer = nn.Sequential(
          nn.Conv2d(1,20,5),
          nn.ReLU(),
          nn.Conv2d(20,64,5),
          nn.ReLU()
        )   

    def forward(self, x):
        output = self.layer(x)
        return output

model = SeqModel()

config = config_from_pytorch_model(model)
output_dir = 'test_pytorch'

convert_from_pytorch_model(
        model, (None, 1, 5, 5), hls_config=config, output_dir=output_dir)

pytests have been added to verify that this keeps working.

Checklist

JanFSchulte commented 1 year ago

pre-commit.ci autofix

vloncar commented 1 year ago

This now includes the changes from #840. There'll be a follow-up PR with the remaining bits we need to parse sPHENIX tracking GNN.