hls4ml fails to read BatchNorm weights from state_dict

[stenv25]

Hello,I want to convert the Pytorch model to HLS and have encountered the following error,I referenced the code in test_merge_pytorch.py:

--------------------------------------------------------------------------
Exception                                 Traceback (most recent call last)
Cell In[14], line 7
      3 config = hls4ml.utils.config_from_pytorch_model(
      4     model, default_precision='ap_fixed<32,16>', inputs_channel_last=True, transpose_outputs=False
      5 )
      6 output_dir = str('model_torch_cnn')
----> 7 hls_model = hls4ml.converters.convert_from_pytorch_model(
      8     model,
      9     [batch_input_shape, batch_input_shape],
     10     hls_config=config,
     11     output_dir=output_dir,
     12     io_type='io_stream',
     13     backend='Vitis',
     14 )
     15 hls_model.compile()

File ~/soft/miniconda3/envs/hls4ml-tutorial/lib/python3.10/site-packages/hls4ml/converters/__init__.py:303, in convert_from_pytorch_model(model, input_shape, output_dir, project_name, input_data_tb, output_data_tb, backend, hls_config, **kwargs)
    299 config['HLSConfig']['Model'] = _check_model_config(model_config)
    301 _check_hls_config(config, hls_config)
--> 303 return pytorch_to_hls(config)

File ~/soft/miniconda3/envs/hls4ml-tutorial/lib/python3.10/site-packages/hls4ml/converters/pytorch_to_hls.py:233, in pytorch_to_hls(config)
    230         raise Exception('Non-default options for groups not implemented yet')
    232 # Process the layer
--> 233 layer, output_shape = layer_handlers[pytorch_class](
    234     pytorch_class, layer_name, input_names, input_shapes, node, class_object, reader, config
    235 )
    237 print(
    238     'Layer name: {}, layer type: {}, input shape: {}'.format(
    239         layer['name'],
   (...)
    242     )
    243 )
    244 layer_list.append(layer)

File ~/soft/miniconda3/envs/hls4ml-tutorial/lib/python3.10/site-packages/hls4ml/converters/pytorch/core.py:111, in parse_batchnorm_layer(operation, layer_name, input_names, input_shapes, node, class_object, data_reader, config)
    108     else:
    109         layer['beta_data'] = 0
--> 111     layer['mean_data'], layer['variance_data'] = get_weights_data(
    112         data_reader, layer['name'], ['running_mean', 'running_variance']
    113     )
    115 in_size = 1
    116 for dim in input_shapes[0][1:]:

File ~/soft/miniconda3/envs/hls4ml-tutorial/lib/python3.10/site-packages/hls4ml/converters/pytorch_to_hls.py:83, in get_weights_data(data_reader, layer_name, var_name)
     80 if not isinstance(var_name, (list, tuple)):
     81     var_name = [var_name]
---> 83 data = [data_reader.get_weights_data(layer_name, var) for var in var_name]
     85 if len(data) == 1:
     86     return data[0]

File ~/soft/miniconda3/envs/hls4ml-tutorial/lib/python3.10/site-packages/hls4ml/converters/pytorch_to_hls.py:83, in <listcomp>(.0)
     80 if not isinstance(var_name, (list, tuple)):
     81     var_name = [var_name]
---> 83 data = [data_reader.get_weights_data(layer_name, var) for var in var_name]
     85 if len(data) == 1:
     86     return data[0]

File ~/soft/miniconda3/envs/hls4ml-tutorial/lib/python3.10/site-packages/hls4ml/converters/pytorch_to_hls.py:39, in PyTorchModelReader.get_weights_data(self, layer_name, var_name)
     36     layer_name = layer_name.split('layer_')[-1]
     38 if var_name not in list(torch_paramap.keys()) + ['weight', 'bias']:
---> 39     raise Exception('Pytorch parameter not yet supported!')
     41 elif var_name in list(torch_paramap.keys()):
     42     var_name = torch_paramap[var_name]

Exception: Pytorch parameter not yet supported!

this is my code:

import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms

transform = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
                                        download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=64,
                                          shuffle=True, num_workers=2)

testset = torchvision.datasets.CIFAR10(root='./data', train=False,
                                       download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=64,
                                         shuffle=False, num_workers=2)

class ConvNet(nn.Module):
    def __init__(self):
        super(ConvNet, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=32, kernel_size=1, stride=1,bias=True)
        self.bn1 = nn.BatchNorm2d(32)
        self.relu1 = nn.ReLU(inplace=True)
        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=1, stride=1,bias=True)
        self.bn2 = nn.BatchNorm2d(64)
        self.relu2 = nn.ReLU(inplace=True)
        self.fc1 = nn.Linear(in_features=64 * 32 * 32, out_features=128)
        self.fc2 = nn.Linear(in_features=128, out_features=10)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu1(x)
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.relu2(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = self.fc2(x)
        return x

model = ConvNet()
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)


def train(net, dataloader, criterion, optimizer):
    net.train()
    running_loss = 0.0
    for i, data in enumerate(dataloader, 0):
        inputs, labels = data
        optimizer.zero_grad()
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
        if i % 100 == 99:
            print('[%d, %5d] loss: %.3f' %
                  (epoch + 1, i + 1, running_loss / 100))
            running_loss = 0.0


def test(net, dataloader, criterion):
    net.eval()
    correct = 0
    total = 0
    loss = 0.0
    with torch.no_grad():
        for data in dataloader:
            images, labels = data
            outputs = net(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
            loss += criterion(outputs, labels).item()
    accuracy = 100 * correct / total
    loss /= len(dataloader)
    print('Accuracy of the network on the test images: %d %%' % (
        accuracy))
    print('Average loss of the network on the test images: %.3f' % (loss))
    return accuracy, loss

istrain = False
if istrain:
    best_acc = 0.0
    for epoch in range(10):
        train(model, trainloader, criterion, optimizer)
        acc, loss = test(model, testloader, criterion)
        if acc > best_acc:
            best_acc = acc
            torch.save(model.state_dict(),'best_model.pt')
else:
    state_dict = torch.load('best_model.pt')
    model.load_state_dict(state_dict,strict=True)
model.eval()




  input_shape = (1,32,32)
  batch_input_shape = (None,) + input_shape
  config = hls4ml.utils.config_from_pytorch_model(
      model, default_precision='ap_fixed<32,16>', inputs_channel_last=True, transpose_outputs=False
  )
  output_dir = str('model_torch_cnn')
  hls_model = hls4ml.converters.convert_from_pytorch_model(
      model,
      [batch_input_shape],
      hls_config=config,
      output_dir=output_dir,
      io_type='io_stream',
      backend='Vitis',
  )
  hls_model.compile()

[vloncar]

With the changes in #840 this model parses.

A few comments:

self.fc1 = nn.Linear(in_features=64 * 32 * 32, out_features=128)

bruh... this layer has over 8 million weights. This won't work in hls4ml.

input_shape = (1,32,32)

The input shape of CIFAR10 is (3, 32, 32)

[stenv25]

Okay, thank you. This is an example written casually.