Empty assertion error

import torch
import torch.nn as nn
import brevitas.nn as qnn
import torch
from torch.utils.data import TensorDataset, DataLoader
import torch.nn as nn
import torch.optim as optim
from tqdm import tqdm
import matplotlib.pyplot as plt

class TinyCNN(nn.Module):
def init(self, n_classes, n_bits) → None:

    a_bits = n_bits
    w_bits = n_bits

    self.q1 = qnn.QuantIdentity(bit_width=a_bits, return_quant_tensor=True)
    self.conv1 = qnn.QuantConv2d(1, 4, 3, stride=1, padding=0, weight_bit_width=w_bits)
    self.q2 = qnn.QuantIdentity(bit_width=a_bits, return_quant_tensor=True)
    self.conv2 = qnn.QuantConv2d(4, 8, 2, stride=2, padding=0, weight_bit_width=w_bits)

    self.fc1 = qnn.QuantLinear(
        8 * 3 * 3,

def forward(self, x):

    x = self.q1(x)
    x = self.conv1(x)
    x = torch.relu(x)
    x = self.q2(x)
    x = self.conv2(x)
    x = torch.relu(x)

    # Flatten the tensor before passing it to the fully connected layer
    x = x.view(x.size(0), -1)

    x = self.fc1(x)
    return x

Now the training part


def train_one_epoch(net, optimizer, train_loader):
# Cross Entropy loss for classification when not using a softmax layer in the network
loss = nn.CrossEntropyLoss()

avg_loss = 0
for data, target in train_loader:
    output = net(data)
    loss_net = loss(output, target.long())
    avg_loss += loss_net.item()

return avg_loss / len(train_loader)

Create a train data loader

train_dataset = TensorDataset(torch.Tensor(x_train), torch.Tensor(y_train))
train_dataloader = DataLoader(train_dataset, batch_size=64, shuffle=True)

test_dataset = TensorDataset(torch.Tensor(x_test), torch.Tensor(y_test))
test_dataloader = DataLoader(test_dataset)

nets = []
bit_range = range(4, 7)

Train the network with Adam, output the test set accuracy every epoch

losses = []
for n_bits in bit_range:
net = TinyCNN(10, n_bits)
losses_bits = []
optimizer = torch.optim.Adam(net.parameters())
for epoch in tqdm(range(N_EPOCHS), desc=f"Training with {n_bits} bit weights and activations"):
losses_bits.append(train_one_epoch(net, optimizer, train_dataloader))



fig = plt.figure(figsize=(8, 4))
for losses_bits in losses:
plt.ylabel(“Cross Entropy Loss”)
plt.legend(list(map(str, bit_range)))
plt.title(“Training set loss during training”)
def test_torch(net, n_bits, test_loader):
all_y_pred = np.zeros(len(test_loader), dtype=np.int64)
all_targets = np.zeros(len(test_loader), dtype=np.int64)

idx = 0
for data, target in test_loader:
    endidx = idx + target.shape[0]
    all_targets[idx:endidx] = target.numpy()

    output = net(data).argmax(1).detach().numpy()
    all_y_pred[idx:endidx] = output

    idx += target.shape[0]

n_correct = np.sum(all_targets == all_y_pred)
print(f"Test accuracy for {n_bits}-bit weights and activations: {n_correct / len(test_loader) * 100:.2f}%")

Test each network in the list

for idx, net in enumerate(nets):
test_torch(net, bit_range[idx], test_dataloader)

try compiling the model

model= compile_brevitas_qat_model(nets[1], x_train,verbose=True)
It is at the point we get an empty assertion error

Hello @Laser_beam , thanks for the report.

Would it be possible to provide (a sub-set of) the data you use ( x_train, x_test, y_train and y_test) in order to be able to reproduce it ? If you prefer, you can probably provide random values but be sure that the types and shapes are matching the original data. Alternatively you could provide a code that generates a similar data-set. Again, I’m mostly interested in the data’s dtypes and shapes.

use the dataset described here

Got it thanks, I’ll take a look at your issue and will get back to you when I know more about it !

1 Like

Hello @Laser_beam,
The problem was indeed not simple to debug but here’s a solution :

  • replace your x = x.view(x.size(0), -1) by x = x.flatten(1) in your forward (works also with x = x.view(-1, 72))
  • add a self.q3 = qnn.QuantIdentity(bit_width=a_bits, return_quant_tensor=True) in your init and a x = self.q3(x) right after your second relu in the forward

The first suggestion is a bit tricky. It is linked to the way we handle shapes under the hood for converting the model to ONNX and generating the underlying FHE circuit. The second one is pretty straightforward : you are missing a quant operator between your ReLU and FC layer ! In fact, if you had used a flatten operator, you would have had a proper error telling you this :sweat_smile:

In any case, we are investigating what could be improved in the code in order to avoid similar confusing errors in the future ! Thanks a lot for the report and keep us updated :wink:

1 Like

Sorry for replying lately , Thankyou @RomanBredehoft
The flattening operation has undergone a transformation – a metamorphosis, . No longer shall we “view” the tensor; instead, we shall have it “flatten” gracefully, :sweat_smile: embracing its true two-dimensional self. :wink:

1 Like