Parsing Flatten layer in PyTorch fails

[stenv25]

Hello, I created a very simple model using PyTorch and found that converting it to HLS failed! Here is my code, can you give me some advice? Thanks very much!

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.flatten = nn.Flatten()
        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.flatten(x)
        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()
print(model)
  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,
      (None, 1,32,32),
      hls_config=config,
      output_dir=output_dir,
      io_type='io_stream',
      backend='Vitis',
  )
  hls_model.compile()

[jmitrevs]

Which version of hls4ml did you use? In the main branch, currently not in a release, the PyTorch support is significantly improved.

[stenv25]

I installed hls4ml under the main branch three weeks ago.

[vloncar]

Fixed in #840, I hope

[stenv25]

OK,Thanks