[ex/mnist] Add a MNIST example

examples/mnist/README.md

@@ -0,0 +1,68 @@
+# LogicNets for MNIST Classification
+
+This example shows the accuracy that is attainable using the LogicNets methodology on the MNIST hand-written character classification task.
+
+## Prerequisites
+
+* LogicNets
+* numpy
+* torchvision
+
+## Installation
+
+If you're using the docker image, all the above prerequisites will be already installed.
+Otherwise, you can install the above dependencies with pip and/or conda.
+
+## Download the Dataset
+
+The MNIST dataset will download automatically when the training script is first run.
+You only need to make sure the necessary directory has been created:
+
+```bash
+mkdir -p data
+```
+
+## Usage
+
+To train the \"MNIST-S\", \"MNIST-M\" and \"MNIST-L\" networks,
+run the following:
+
+```bash
+python train.py --arch <mnist-s|mnist-m|mnist-l> --log-dir ./<mnist_s|mnist_m|mnist_l>/
+```
+
+To then generate verilog from this trained model, run the following:
+
+```bash
+python neq2lut.py --arch <mnist-s|mnist-m|mnist-l> --checkpoint ./<mnist_s|mnist_m|mnist_l>/best_accuracy.pth --log-dir ./<mnist_s|mnist_m|mnist_l>/verilog/ --add-registers
+```
+
+## Results
+
+Your results may vary slightly, depending on your system configuration.
+The following results are attained when training on a CPU and synthesising with Vivado 2019.2:
+
+| Network Architecture  | Test Accuracy (%) | LUTs  | Flip Flops    | Fmax (Mhz)    | Latency (Cycles)  |
+| --------------------- | ----------------- | ----- | ------------- | ------------- | ----------------- |
+| MNIST-S               |                   |       |               |               |                   |
+| MNIST-M               |                   |       |               |               |                   |
+| MNIST-L               |                   |       |               |               |                   |
+
+## Citation
+
+If you find this work useful for your research, please consider citing
+our paper below:
+
+```bibtex
+@inproceedings{umuroglu2020logicnets,
+  author = {Umuroglu, Yaman and Akhauri, Yash and Fraser, Nicholas J and Blott, Michaela},
+  booktitle = {Proceedings of the International Conference on Field-Programmable Logic and Applications},
+  title = {LogicNets: Co-Designed Neural Networks and Circuits for Extreme-Throughput Applications},
+  year = {2020},
+  pages = {291-297},
+  publisher = {IEEE Computer Society},
+  address = {Los Alamitos, CA, USA},
+  month = {sep}
+}
+```
+

examples/mnist/dataset_dump.py

@@ -0,0 +1,131 @@
+#  Copyright (C) 2021 Xilinx, Inc
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+import os
+from argparse import ArgumentParser
+from functools import reduce, partial
+
+import torch
+from torch.utils.data import DataLoader
+from torchvision.datasets import MNIST
+from torchvision import transforms
+
+from logicnets.nn import    generate_truth_tables, \
+                            lut_inference, \
+                            module_list_to_verilog_module
+from logicnets.synthesis import synthesize_and_get_resource_counts
+
+from train import configs, model_config, other_options, test
+from models import MnistNeqModel, MnistLutModel
+
+def dump_io(model, data_loader, input_file, output_file):
+    input_quant = model.module_list[0].input_quant
+    _, input_bitwidth = input_quant.get_scale_factor_bits()
+    input_bitwidth = int(input_bitwidth)
+    total_input_bits = model.module_list[0].in_features*input_bitwidth
+    input_quant.bin_output()
+    with open(input_file, 'w') as i_f, open(output_file, 'w') as o_f:
+        for data, target in data_loader:
+            x = input_quant(data)
+            indices = target
+            for i in range(x.shape[0]):
+                x_i = x[i,:]
+                xv_i = list(map(lambda z: input_quant.get_bin_str(z), x_i))
+                xvc_i = reduce(lambda a,b: a+b, xv_i[::-1])
+                i_f.write(f"{int(xvc_i,2):0{int(total_input_bits)}b}\n")
+                o_f.write(f"{int(indices[i])}\n")
+
+if __name__ == "__main__":
+    parser = ArgumentParser(description="Dump the train and test datasets (after input quantization) into text files")
+    parser.add_argument('--arch', type=str, choices=configs.keys(), default="mnist-s",
+        help="Specific the neural network model to use (default: %(default)s)")
+    parser.add_argument('--batch-size', type=int, default=None, metavar='N',
+        help="Batch size for evaluation (default: %(default)s)")
+    parser.add_argument('--input-bitwidth', type=int, default=None,
+        help="Bitwidth to use at the input (default: %(default)s)")
+    parser.add_argument('--hidden-bitwidth', type=int, default=None,
+        help="Bitwidth to use for activations in hidden layers (default: %(default)s)")
+    parser.add_argument('--output-bitwidth', type=int, default=None,
+        help="Bitwidth to use at the output (default: %(default)s)")
+    parser.add_argument('--input-fanin', type=int, default=None,
+        help="Fanin to use at the input (default: %(default)s)")
+    parser.add_argument('--hidden-fanin', type=int, default=None,
+        help="Fanin to use for the hidden layers (default: %(default)s)")
+    parser.add_argument('--output-fanin', type=int, default=None,
+        help="Fanin to use at the output (default: %(default)s)")
+    parser.add_argument('--hidden-layers', nargs='+', type=int, default=None,
+        help="A list of hidden layer neuron sizes (default: %(default)s)")
+    parser.add_argument('--log-dir', type=str, default='./log',
+        help="A location to store the output I/O text files (default: %(default)s)")
+    parser.add_argument('--checkpoint', type=str, required=True,
+        help="The checkpoint file which contains the model weights")
+    args = parser.parse_args()
+    defaults = configs[args.arch]
+    options = vars(args)
+    del options['arch']
+    config = {}
+    for k in options.keys():
+        config[k] = options[k] if options[k] is not None else defaults[k] # Override defaults, if specified.
+
+    if not os.path.exists(config['log_dir']):
+        os.makedirs(config['log_dir'])
+
+    # Split up configuration options to be more understandable
+    model_cfg = {}
+    for k in model_config.keys():
+        model_cfg[k] = config[k]
+    options_cfg = {}
+    for k in other_options.keys():
+        if k == 'cuda':
+            continue
+        options_cfg[k] = config[k]
+
+    trans = transform=transforms.Compose([
+                transforms.ToTensor(),
+                transforms.Normalize((0.1307,), (0.3081,)),
+                transforms.Lambda(partial(torch.reshape, shape=(-1,)))
+            ])
+
+    # Fetch the datasets
+    dataset = {}
+    dataset['train'] = MNIST('./data', train=True, download=True, transform=trans)
+    dataset['test'] = MNIST('./data', train=False, download=True, transform=trans)
+    train_loader = DataLoader(dataset["train"], batch_size=config['batch_size'], shuffle=False)
+    test_loader = DataLoader(dataset["test"], batch_size=config['batch_size'], shuffle=False)
+
+    # Instantiate the PyTorch model
+    x, y = dataset["train"][0]
+    model_cfg['input_length'] = len(x)
+    model_cfg['output_length'] = 10
+    model = MnistNeqModel(model_cfg)
+
+    # Load the model weights
+    checkpoint = torch.load(options_cfg['checkpoint'], map_location='cpu')
+    model.load_state_dict(checkpoint['model_dict'])
+
+    # Test the PyTorch model
+    print("Running inference on baseline model...")
+    model.eval()
+    baseline_accuracy = test(model, test_loader, cuda=False)
+    print("Baseline accuracy: %f" % (baseline_accuracy))
+
+    # Run preprocessing on training set.
+    train_input_file = config['log_dir'] + "/train_input.txt"
+    train_output_file = config['log_dir'] + "/train_output.txt"
+    test_input_file = config['log_dir'] + "/test_input.txt"
+    test_output_file = config['log_dir'] + "/test_output.txt"
+    print(f"Dumping train I/O to {train_input_file} and {train_output_file}")
+    dump_io(model, train_loader, train_input_file, train_output_file)
+    print(f"Dumping test I/O to {test_input_file} and {test_output_file}")
+    dump_io(model, test_loader, test_input_file, test_output_file)

examples/mnist/models.py

@@ -0,0 +1,145 @@
+#  Copyright (C) 2021 Xilinx, Inc
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+from functools import reduce
+from os.path import realpath
+
+import torch
+import torch.nn as nn
+from torch.nn.parameter import Parameter
+from torch.nn import init
+
+from brevitas.core.quant import QuantType
+from brevitas.core.scaling import ScalingImplType
+from brevitas.nn import QuantHardTanh, QuantReLU
+
+from pyverilator import PyVerilator
+
+from logicnets.quant import QuantBrevitasActivation
+from logicnets.nn import SparseLinearNeq, ScalarBiasScale, RandomFixedSparsityMask2D, DenseMask2D
+from logicnets.init import random_restrict_fanin
+
+class MnistNeqModel(nn.Module):
+    def __init__(self, model_config):
+        super(MnistNeqModel, self).__init__()
+        self.model_config = model_config
+        self.num_neurons = [model_config["input_length"]] + model_config["hidden_layers"] + [model_config["output_length"]]
+        layer_list = []
+        for i in range(1, len(self.num_neurons)):
+            in_features = self.num_neurons[i-1]
+            out_features = self.num_neurons[i]
+            bn = nn.BatchNorm1d(out_features)
+            nn.init.constant_(bn.weight.data, 1)
+            nn.init.constant_(bn.bias.data, 0)
+            if i == 1:
+                do_in = nn.Dropout(p=model_config["input_dropout"])
+                bn_in = nn.BatchNorm1d(in_features)
+                nn.init.constant_(bn_in.weight.data, 1)
+                nn.init.constant_(bn_in.bias.data, 0)
+                input_bias = ScalarBiasScale(scale=False, bias_init=-0.25)
+                input_quant = QuantBrevitasActivation(QuantHardTanh(model_config["input_bitwidth"], max_val=1., narrow_range=False, quant_type=QuantType.INT, scaling_impl_type=ScalingImplType.PARAMETER), pre_transforms=[do_in, bn_in, input_bias])
+                output_quant = QuantBrevitasActivation(QuantReLU(bit_width=model_config["hidden_bitwidth"], max_val=1.61, quant_type=QuantType.INT, scaling_impl_type=ScalingImplType.PARAMETER), pre_transforms=[bn])
+                mask = RandomFixedSparsityMask2D(in_features, out_features, fan_in=model_config["input_fanin"])
+                layer = SparseLinearNeq(in_features, out_features, input_quant=input_quant, output_quant=output_quant, sparse_linear_kws={'mask': mask})
+                layer_list.append(layer)
+            elif i == len(self.num_neurons)-1:
+                output_bias_scale = ScalarBiasScale(bias_init=0.33)
+                output_quant = QuantBrevitasActivation(QuantHardTanh(bit_width=model_config["output_bitwidth"], max_val=1.33, narrow_range=False, quant_type=QuantType.INT, scaling_impl_type=ScalingImplType.PARAMETER), pre_transforms=[bn], post_transforms=[output_bias_scale])
+                mask = RandomFixedSparsityMask2D(in_features, out_features, fan_in=model_config["output_fanin"])
+                layer = SparseLinearNeq(in_features, out_features, input_quant=layer_list[-1].output_quant, output_quant=output_quant, sparse_linear_kws={'mask': mask}, apply_input_quant=False)
+                layer_list.append(layer)
+            else:
+                output_quant = QuantBrevitasActivation(QuantReLU(bit_width=model_config["hidden_bitwidth"], max_val=1.61, quant_type=QuantType.INT, scaling_impl_type=ScalingImplType.PARAMETER), pre_transforms=[bn])
+                mask = RandomFixedSparsityMask2D(in_features, out_features, fan_in=model_config["hidden_fanin"])
+                layer = SparseLinearNeq(in_features, out_features, input_quant=layer_list[-1].output_quant, output_quant=output_quant, sparse_linear_kws={'mask': mask}, apply_input_quant=False)
+                layer_list.append(layer)
+        self.module_list = nn.ModuleList(layer_list)
+        self.is_verilog_inference = False
+        self.latency = 1
+        self.verilog_dir = None
+        self.top_module_filename = None
+        self.dut = None
+        self.logfile = None
+
+    def verilog_inference(self, verilog_dir, top_module_filename, logfile: bool = False, add_registers: bool = False):
+        self.verilog_dir = realpath(verilog_dir)
+        self.top_module_filename = top_module_filename
+        self.dut = PyVerilator.build(f"{self.verilog_dir}/{self.top_module_filename}", verilog_path=[self.verilog_dir], build_dir=f"{self.verilog_dir}/verilator")
+        self.is_verilog_inference = True
+        self.logfile = logfile
+        if add_registers:
+            self.latency = len(self.num_neurons)
+
+    def pytorch_inference(self):
+        self.is_verilog_inference = False
+
+    def verilog_forward(self, x):
+        # Get integer output from the first layer
+        input_quant = self.module_list[0].input_quant
+        output_quant = self.module_list[-1].output_quant
+        _, input_bitwidth = self.module_list[0].input_quant.get_scale_factor_bits()
+        _, output_bitwidth = self.module_list[-1].output_quant.get_scale_factor_bits()
+        input_bitwidth, output_bitwidth = int(input_bitwidth), int(output_bitwidth)
+        total_input_bits = self.module_list[0].in_features*input_bitwidth
+        total_output_bits = self.module_list[-1].out_features*output_bitwidth
+        num_layers = len(self.module_list)
+        input_quant.bin_output()
+        self.module_list[0].apply_input_quant = False
+        y = torch.zeros(x.shape[0], self.module_list[-1].out_features)
+        x = input_quant(x)
+        self.dut.io.rst = 0
+        self.dut.io.clk = 0
+        for i in range(x.shape[0]):
+            x_i = x[i,:]
+            y_i = self.pytorch_forward(x[i:i+1,:])[0]
+            xv_i = list(map(lambda z: input_quant.get_bin_str(z), x_i))
+            ys_i = list(map(lambda z: output_quant.get_bin_str(z), y_i))
+            xvc_i = reduce(lambda a,b: a+b, xv_i[::-1])
+            ysc_i = reduce(lambda a,b: a+b, ys_i[::-1])
+            self.dut["M0"] = int(xvc_i, 2)
+            for j in range(self.latency + 1):
+                #print(self.dut.io.M5)
+                res = self.dut[f"M{num_layers}"]
+                result = f"{res:0{int(total_output_bits)}b}"
+                self.dut.io.clk = 1
+                self.dut.io.clk = 0
+            expected = f"{int(ysc_i,2):0{int(total_output_bits)}b}"
+            result = f"{res:0{int(total_output_bits)}b}"
+            assert(expected == result)
+            res_split = [result[i:i+output_bitwidth] for i in range(0, len(result), output_bitwidth)][::-1]
+            yv_i = torch.Tensor(list(map(lambda z: int(z, 2), res_split)))
+            y[i,:] = yv_i
+            # Dump the I/O pairs
+            if self.logfile is not None:
+                with open(self.logfile, "a") as f:
+                    f.write(f"{int(xvc_i,2):0{int(total_input_bits)}b}{int(ysc_i,2):0{int(total_output_bits)}b}\n")
+        return y
+
+    def pytorch_forward(self, x):
+        for l in self.module_list:
+            x = l(x)
+        return x
+
+    def forward(self, x):
+        if self.is_verilog_inference:
+            return self.verilog_forward(x)
+        else:
+            return self.pytorch_forward(x)
+
+class MnistLutModel(MnistNeqModel):
+    pass
+
+class MnistVerilogModel(MnistNeqModel):
+    pass
+