mnist_mlp/gen_from_model.py

+import argparse
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from torchvision import datasets, transforms
+from torch.optim.lr_scheduler import StepLR
+from gen_weight import gen_weight, gen_input
+from create_config_pkg import *
+
+def main():
+    # Training settings
+    # read the help to know which parameter do what
+    # type $ python3 main.py --help
+    parser = argparse.ArgumentParser(description='script for exporting weights, inputs and config from PyTorch models.')
+    parser.add_argument('--model-path', type=str, default="./mon_model.mpt",
+                        help='indication suffix for your model name')
+
+    parser.add_argument('--export-weights', action='store_true', default=False,
+                        help='export weights in hexadecimal for hardware implementation') # export weights in files for HW implementation
+
+    parser.add_argument('--export-inputs', action='store_true', default=False,
+                        help='export inputs in hexadecimal for hardware implementation') # export inputs in .mif files for hardware test
+
+    parser.add_argument('--num-of-sample', type=int, default=16,
+                        help='number of input samples exported') 
+
+    parser.add_argument('--export-path', type=str, default="./data_files",
+                        help='path for stroring the exported weights, inputs and config files, (default = ./data_files)') #path 
+
+    parser.add_argument('--bit-width', type=int, default=32,
+                        help='number of bits per weights and inputs') 
+
+    parser.add_argument('--int-width', type=int, default=7,
+                        help='size of the integer part in the fixe points weights') 
+
+    parser.add_argument('--gen-config', action='store_true', default=False,
+                        help='generate a VHDL package for configuration purposes for your hardware implementation') 
+
+    parser.add_argument('--config-name', type=str, default="config_test",
+                        help='name of the config file, (default = config_test)')
+
+
+    if args.export_weights :
+        gen_weight(model, bit_width = args.bit_width, int_width = args.int_width, den_width = args.bit_width-args.int_width ,dirname = args.export_path) # gen quantized weights for hw implemenetation
+    
+    if args.export_inputs :
+        gen_input(test_loader, num_samples = args.num_of_sample, bit_width = args.bit_width, int_width = args.int_width, den_width = args.bit_width-args.int_width ,dirname = args.export_path) # gen inputs quantized as well
+
+    if args.gen_config :
+        gen_config_pkg(model, bit_width = args.bit_width, int_width = args.bit_width, dirname = args.export_path, config_name = args.config_name)
+
+if __name__ == '__main__':
+    main()

mnist_mlp/gen_weight.py

+import numpy as np
+import argparse
+import os
+import csv
+import torch.nn as nn
+import torch
+
+export_path = "./"
+default_dirname = "default"
+
+def DtoB(num, bit_width, int_width, den_width):   # Decimal to binary converter
+    # saturations
+    if num > 2**(int_width-1) : # saturate over the max representable value
+        num = 2**(int_width-1)
+    elif num < -2**(int_width-1) : # saturate under the max negative representable value
+        num = -2**(int_width-1)
+ 
+    num = num * (2**den_width) # shift left of den_width
+    if num < 0 :
+        # if num is negative => 2's complement
+        num = -num
+        b = int(round(num)) # conversion in int
+        b = ~b
+        d = b+1
+    elif num == 0 :
+        d = 0
+    else :
+        d = int(round(num)) # conversion in int
+    return d
+
+def gen_weight(model,bit_width = 32, int_width = 7, den_width = 25,dirname = default_dirname):
+    # generate a csv file for each neurons with quantized weights in fixe point (note that you can choose which part represent the fraction and integer part)
+
+    # model => model you want to exctract weights from
+    # bit_width => bit width of the wanted data (usualy 8, 16, 32, 64 and so on), note that pytorch store values in fp32 for floating point on 32 bits (see : https://en.wikipedia.org/wiki/IEEE_754)
+    # int_width => bit width of the integer part of the fixe point representation (/!\ (2^int_width)/2 = max int) remember the first bit is used as a sign bit
+    # den_width => bit width of the fraction part, (smallest representable number = 2^-den_width)
+    layers = list(model.children())
+    
+    if not os.path.exists(dirname):
+        os.makedirs(dirname)
+
+    subdir_weights = os.path.join(dirname, "weights") 
+    if not os.path.exists(subdir_weights):
+        os.makedirs(subdir_weights) # makes a subdirectory to store weights
+    else :
+        print("warning : weights files already generate")
+
+    cnt_l = 0
+    for layer in layers:  # Exclude the last layer
+        if isinstance(layer, nn.Linear):
+            # Get the number of rows and columns
+            weights_tensor = layer.weight
+            biases_tensor = layer.bias
+            with torch.no_grad():
+                weights_tensor = weights_tensor.cpu().numpy()
+                biases_tensor = biases_tensor.cpu().numpy()
+            rows, cols = weights_tensor.shape
+
+            for row in range(rows) : # rows represent neurons in the pytorch way
+                neuron_weights = weights_tensor[row,:]
+                neuron_biase = biases_tensor[row]
+                # Open the file in write mode
+                # print('layer', cnt+1)
+                # print('neuron', row+1)
+                print("cnt",cnt_l)
+                csv_path = os.path.join(subdir_weights, f'w_l{cnt_l}_n{row}.mif')
+                with open(csv_path, 'w', newline='') as file:
+                    # Convert each integer to hexadecimal and write one per line
+                    for val in neuron_weights:
+                        val = DtoB(val, bit_width, int_width, den_width) 
+                        if bin(val)[0] == "-" :
+                            file.write(bin(val)[3:] + '\n')
+                        else :
+                            file.write(bin(val)[2:] + '\n')
+                    neuron_biase = DtoB(neuron_biase, bit_width, int_width, den_width) 
+                    if bin(neuron_biase)[0] == "-" :
+                        file.write(bin(neuron_biase)[3:] + '\n')
+                    else :
+                        file.write(bin(neuron_biase)[2:] + '\n')
+            cnt_l += 1
+
+def gen_input(dataloader, num_samples = 10, bit_width = 32, int_width = 7, den_width = 25, dirname = default_dirname):
+    # Create the output directory if it doesn't exist
+    if not os.path.exists(dirname):
+        os.makedirs(dirname)
+
+    subdir_inputs = os.path.join(dirname, "inputs") 
+    if not os.path.exists(subdir_inputs):
+        os.makedirs(subdir_inputs) # makes a subdirectory to store weights
+
+
+
+    for batch_idx, (image, label) in enumerate(dataloader): # shuffle the dataset otherwise it won't be random
+        if batch_idx < 1 :
+            print(batch_idx)
+            with torch.no_grad():
+                image_array = image.cpu().numpy()
+                label_array = label.cpu().numpy() 
+            batch_dim = image_array.shape[0]
+            cnt = 0
+            for image in range(batch_dim) :
+                if image <= num_samples :
+                    csv_filename = os.path.join(subdir_inputs, f'in_{cnt}.mif')
+                    with open(csv_filename, 'w', newline='') as file:
+                        # Write the pixel values in hexadecimal format, one per line
+                        for pixel in image_array[image, :, :].flatten():
+                            pixel = DtoB(pixel, bit_width, int_width, den_width) 
+                            file.write(bin(pixel)[2:]+ '\n')
+                else :
+                    break
+                cnt += 1
+        else :
+            break
+
+
+def find_max(model, verbose = False) : # find max and min weight of the model for selectionning int_width
+    max_w = float('-inf')
+    max_b = float('-inf')
+    min_w = float('inf')
+    min_b = float('inf')
+    layers = list(model.children())
+    for layer in layers : 
+        if isinstance(layer, nn.Linear):
+            # Get the number of rows and columns
+            weights_tensor = layer.weight
+            biases_tensor = layer.bias
+            with torch.no_grad():
+                weights_tensor = weights_tensor.numpy()
+                biases_tensor = biases_tensor.numpy()
+
+            max_w = max([max_w, np.max(weights_tensor)])
+            max_n = max([max_b, np.max(biases_tensor)])
+            max_t = max([max_w, max_b])
+            if verbose :print(f"max : {max_t}")
+
+            min_w = min([min_w, np.min(weights_tensor)])
+            min_n = min([min_b, np.min(biases_tensor)])
+            min_t = min([min_w, min_b])
+            if verbose : print(f"min : {min_t}")
+    return min_t, max_t
+
+def print_model(model) : # print all the layers of the model
+    layers = list(model.children())
+    for layer in layers : 
+        if isinstance(layer, nn.Linear):
+            print(layer)

mnist_mlp/lib64

+lib
\ No newline at end of file

mnist_mlp/main.py

+import argparse
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from torchvision import datasets, transforms
+from torch.optim.lr_scheduler import StepLR
+from gen_weight import gen_weight, gen_input
+from create_config_pkg import *
+
+# NETWORK DEFINITON
+# here you can change the structure of the network
+# /!\ mnist input = 28*28 digit pictures, hence we get 784 pixels
+class Net(nn.Module):
+    def __init__(self):
+        super(Net, self).__init__()
+        self.fc1 = nn.Linear(784, 16) # achrtung, flattening needed
+        self.hardtanhfc1 = nn.Hardtanh()  # Activation function after fc1
+        self.fc2 = nn.Linear(16, 16) 
+        self.hardtanhfc2 = nn.Hardtanh()  # Activation function after fc2
+        self.fc3 = nn.Linear(16, 10)
+
+    def forward(self, x):
+        x = torch.flatten(x, 1)
+        x = self.fc1(x)
+        x = self.hardtanhfc1(x) # here you can change the activation function between fc1 and fc2
+        x = self.fc2(x)
+        x =  self.hardtanhfc2(x)
+        output = self.fc3(x)
+        return output
+
+
+def train(args, model, device, train_loader, optimizer, epoch, criterion):
+    model.train()
+    for batch_idx, (data, target) in enumerate(train_loader):
+        data, target = data.to(device), target.to(device)
+        data = data
+        optimizer.zero_grad()
+        output = model(data)
+        loss = criterion(output, target) # negative log likelihood loss 
+        loss.backward()
+        optimizer.step()
+        if batch_idx % args.log_interval == 0:
+            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
+                epoch, batch_idx * len(data), len(train_loader.dataset),
+                100. * batch_idx / len(train_loader), loss.item()))
+            if args.dry_run:
+                break
+
+
+def test(model, device, test_loader, criterion):
+    model.eval()
+    test_loss = 0
+    correct = 0
+    with torch.no_grad():
+        for data, target in test_loader:
+            data = data
+            data, target = data.to(device), target.to(device)
+            output = model(data)
+            test_loss += criterion(output, target)  # sum up batch loss
+            pred = output.argmax(dim=1, keepdim=True)  # get the index of the max output (hardmax)
+            correct += pred.eq(target.view_as(pred)).sum().item()
+
+    test_loss /= len(test_loader.dataset)
+
+    print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
+        test_loss, correct, len(test_loader.dataset),
+        100. * correct / len(test_loader.dataset)))
+
+
+def main():
+    # Training settings
+    # read the help to know which parameter do what
+    # type $ python3 main.py --help
+    parser = argparse.ArgumentParser(description='PyTorch MNIST Example, and hardware implementation')
+    parser.add_argument('--batch-size', type=int, default=64, metavar='N',
+                        help='input batch size for training (default: 64)') # train batch size
+
+    parser.add_argument('--test-batch-size', type=int, default=1000, metavar='N',
+                        help='input batch size for testing (default: 1000)') # test batch size
+
+    parser.add_argument('--epochs', type=int, default=13, metavar='N',
+                        help='number of epochs to train (default: 13)') # number of epoch
+
+    parser.add_argument('--lr', type=float, default=0.001, metavar='LR',
+                        help='learning rate (default: 0.001)') # learning rate
+
+    parser.add_argument('--gamma', type=float, default=0.7, metavar='M',
+                        help='Learning rate step gamma (default: 0.7)') # 
+
+    parser.add_argument('--no-cuda', action='store_true', default=False,
+                        help='disables CUDA training') # enable cuda (GPU training)
+
+    parser.add_argument('--no-mps', action='store_true', default=False,
+                        help='disables macOS GPU training') # enable GPU training for macOS
+
+    parser.add_argument('--dry-run', action='store_true', default=False,
+                        help='quickly check a single pass') # dry run, no iteration just for test
+
+    parser.add_argument('--seed', type=int, default=23, metavar='S',
+                        help='random seed (default: 23)') # random seed for reproductible results
+
+    parser.add_argument('--log-interval', type=int, default=10, metavar='N',
+                        help='how many batches to wait before logging training status')
+
+    parser.add_argument('--save-model', action='store_true', default=False,
+                        help='For Saving the current Model')
+
+    parser.add_argument('--model-suffix', type=str, default="mon_model",
+                        help='indication suffix for your model name')
+
+    parser.add_argument('--export-weights', action='store_true', default=False,
+                        help='export weights in hexadecimal for hardware implementation') # export weights in files for HW implementation
+
+    parser.add_argument('--export-inputs', action='store_true', default=False,
+                        help='export inputs in hexadecimal for hardware implementation') # export inputs in .mif files for hardware test
+
+    parser.add_argument('--num-of-sample', type=int, default=16,
+                        help='number of input samples exported') 
+
+    parser.add_argument('--export-path', type=str, default="./data_files",
+                        help='path for stroring the exported weights, inputs and config files, (default = ./data_files)') #path 
+
+    parser.add_argument('--bit-width', type=int, default=32,
+                        help='number of bits per weights and inputs') 
+
+    parser.add_argument('--int-width', type=int, default=7,
+                        help='size of the integer part in the fixe points weights') 
+
+    parser.add_argument('--gen-config', action='store_true', default=False,
+                        help='generate a VHDL package for configuration purposes for your hardware implementation') 
+
+    parser.add_argument('--config-name', type=str, default="config_test",
+                        help='name of the config file, (default = config_test)')
+
+    args = parser.parse_args()
+    use_cuda = not args.no_cuda and torch.cuda.is_available()
+    use_mps = not args.no_mps and torch.backends.mps.is_available()
+
+    torch.manual_seed(args.seed)
+
+    if use_cuda:
+        device = torch.device("cuda")
+    elif use_mps:
+        device = torch.device("mps")
+    else:
+        device = torch.device("cpu")
+
+    train_kwargs = {'batch_size': args.batch_size}
+    test_kwargs = {'batch_size': args.test_batch_size}
+    if use_cuda:
+        cuda_kwargs = {'num_workers': 1,
+                       'pin_memory': True,
+                       'shuffle': True}
+        train_kwargs.update(cuda_kwargs)
+        test_kwargs.update(cuda_kwargs)
+
+    transform=transforms.Compose([
+        transforms.ToTensor(),
+        ])
+    dataset1 = datasets.MNIST('../data', train=True, download=True,
+                       transform=transform)
+    dataset2 = datasets.MNIST('../data', train=False,
+                       transform=transform)
+    train_loader = torch.utils.data.DataLoader(dataset1,**train_kwargs)
+    test_loader = torch.utils.data.DataLoader(dataset2, **test_kwargs)
+
+    model = Net().to(device)
+    criterion = nn.CrossEntropyLoss()
+    optimizer = torch.optim.Adam(model.parameters())
+
+    scheduler = StepLR(optimizer, step_size=1, gamma=args.gamma)
+    for epoch in range(1, args.epochs + 1):
+        train(args, model, device, train_loader, optimizer, epoch, criterion)
+        test(model, device, test_loader, criterion)
+        scheduler.step()
+
+    if args.save_model:
+        torch.save(model.state_dict(), f"mnist_mlp_{args.model_suffix}.pt")
+
+    if args.export_weights :
+        gen_weight(model, bit_width = args.bit_width, int_width = args.int_width, den_width = args.bit_width-args.int_width ,dirname = args.export_path) # gen quantized weights for hw implemenetation
+    
+    if args.export_inputs :
+        gen_input(test_loader, num_samples = args.num_of_sample, bit_width = args.bit_width, int_width = args.int_width, den_width = args.bit_width-args.int_width ,dirname = args.export_path) # gen inputs quantized as well
+
+    if args.gen_config :
+        gen_config_pkg(model, bit_width = args.bit_width, int_width = args.int_width, dirname = args.export_path, config_name = args.config_name)
+
+if __name__ == '__main__':
+    main()

mnist_mlp/pyvenv.cfg

+home = /usr/bin
+include-system-site-packages = false
+version = 3.10.12

mnist_mlp/share/man/man1/isympy.1

+'\" -*- coding: us-ascii -*-
+.if \n(.g .ds T< \\FC
+.if \n(.g .ds T> \\F[\n[.fam]]
+.de URL
+\\$2 \(la\\$1\(ra\\$3
+..
+.if \n(.g .mso www.tmac
+.TH isympy 1 2007-10-8 "" ""
+.SH NAME
+isympy \- interactive shell for SymPy
+.SH SYNOPSIS
+'nh
+.fi
+.ad l
+\fBisympy\fR \kx
+.if (\nx>(\n(.l/2)) .nr x (\n(.l/5)
+'in \n(.iu+\nxu
+[\fB-c\fR | \fB--console\fR] [\fB-p\fR ENCODING | \fB--pretty\fR ENCODING] [\fB-t\fR TYPE | \fB--types\fR TYPE] [\fB-o\fR ORDER | \fB--order\fR ORDER] [\fB-q\fR | \fB--quiet\fR] [\fB-d\fR | \fB--doctest\fR] [\fB-C\fR | \fB--no-cache\fR] [\fB-a\fR | \fB--auto\fR] [\fB-D\fR | \fB--debug\fR] [
+-- | PYTHONOPTIONS]
+'in \n(.iu-\nxu
+.ad b
+'hy
+'nh
+.fi
+.ad l
+\fBisympy\fR \kx
+.if (\nx>(\n(.l/2)) .nr x (\n(.l/5)
+'in \n(.iu+\nxu
+[
+{\fB-h\fR | \fB--help\fR}
+|
+{\fB-v\fR | \fB--version\fR}
+]
+'in \n(.iu-\nxu
+.ad b
+'hy
+.SH DESCRIPTION
+isympy is a Python shell for SymPy. It is just a normal python shell
+(ipython shell if you have the ipython package installed) that executes
+the following commands so that you don't have to:
+.PP
+.nf
+\*(T<
+>>> from __future__ import division
+>>> from sympy import *
+>>> x, y, z = symbols("x,y,z")
+>>> k, m, n = symbols("k,m,n", integer=True)
+    \*(T>
+.fi
+.PP
+So starting isympy is equivalent to starting python (or ipython) and
+executing the above commands by hand. It is intended for easy and quick
+experimentation with SymPy. For more complicated programs, it is recommended
+to write a script and import things explicitly (using the "from sympy
+import sin, log, Symbol, ..." idiom).
+.SH OPTIONS
+.TP
+\*(T<\fB\-c \fR\*(T>\fISHELL\fR, \*(T<\fB\-\-console=\fR\*(T>\fISHELL\fR
+Use the specified shell (python or ipython) as
+console backend instead of the default one (ipython
+if present or python otherwise).
+
+Example: isympy -c python
+
+\fISHELL\fR could be either
+\&'ipython' or 'python'
+.TP
+\*(T<\fB\-p \fR\*(T>\fIENCODING\fR, \*(T<\fB\-\-pretty=\fR\*(T>\fIENCODING\fR
+Setup pretty printing in SymPy. By default, the most pretty, unicode
+printing is enabled (if the terminal supports it). You can use less
+pretty ASCII printing instead or no pretty printing at all.
+
+Example: isympy -p no
+
+\fIENCODING\fR must be one of 'unicode',
+\&'ascii' or 'no'.
+.TP
+\*(T<\fB\-t \fR\*(T>\fITYPE\fR, \*(T<\fB\-\-types=\fR\*(T>\fITYPE\fR
+Setup the ground types for the polys. By default, gmpy ground types
+are used if gmpy2 or gmpy is installed, otherwise it falls back to python
+ground types, which are a little bit slower. You can manually
+choose python ground types even if gmpy is installed (e.g., for testing purposes).
+
+Note that sympy ground types are not supported, and should be used
+only for experimental purposes.
+
+Note that the gmpy1 ground type is primarily intended for testing; it the
+use of gmpy even if gmpy2 is available.
+
+This is the same as setting the environment variable
+SYMPY_GROUND_TYPES to the given ground type (e.g.,
+SYMPY_GROUND_TYPES='gmpy')
+
+The ground types can be determined interactively from the variable
+sympy.polys.domains.GROUND_TYPES inside the isympy shell itself.
+
+Example: isympy -t python
+
+\fITYPE\fR must be one of 'gmpy',
+\&'gmpy1' or 'python'.
+.TP
+\*(T<\fB\-o \fR\*(T>\fIORDER\fR, \*(T<\fB\-\-order=\fR\*(T>\fIORDER\fR
+Setup the ordering of terms for printing. The default is lex, which
+orders terms lexicographically (e.g., x**2 + x + 1). You can choose
+other orderings, such as rev-lex, which will use reverse
+lexicographic ordering (e.g., 1 + x + x**2).
+
+Note that for very large expressions, ORDER='none' may speed up
+printing considerably, with the tradeoff that the order of the terms
+in the printed expression will have no canonical order
+
+Example: isympy -o rev-lax
+
+\fIORDER\fR must be one of 'lex', 'rev-lex', 'grlex',
+\&'rev-grlex', 'grevlex', 'rev-grevlex', 'old', or 'none'.
+.TP
+\*(T<\fB\-q\fR\*(T>, \*(T<\fB\-\-quiet\fR\*(T>
+Print only Python's and SymPy's versions to stdout at startup, and nothing else.
+.TP
+\*(T<\fB\-d\fR\*(T>, \*(T<\fB\-\-doctest\fR\*(T>
+Use the same format that should be used for doctests. This is
+equivalent to '\fIisympy -c python -p no\fR'.
+.TP
+\*(T<\fB\-C\fR\*(T>, \*(T<\fB\-\-no\-cache\fR\*(T>
+Disable the caching mechanism. Disabling the cache may slow certain
+operations down considerably. This is useful for testing the cache,
+or for benchmarking, as the cache can result in deceptive benchmark timings.
+
+This is the same as setting the environment variable SYMPY_USE_CACHE
+to 'no'.
+.TP
+\*(T<\fB\-a\fR\*(T>, \*(T<\fB\-\-auto\fR\*(T>
+Automatically create missing symbols. Normally, typing a name of a
+Symbol that has not been instantiated first would raise NameError,
+but with this option enabled, any undefined name will be
+automatically created as a Symbol. This only works in IPython 0.11.
+
+Note that this is intended only for interactive, calculator style
+usage. In a script that uses SymPy, Symbols should be instantiated
+at the top, so that it's clear what they are.
+
+This will not override any names that are already defined, which
+includes the single character letters represented by the mnemonic
+QCOSINE (see the "Gotchas and Pitfalls" document in the
+documentation). You can delete existing names by executing "del
+name" in the shell itself. You can see if a name is defined by typing
+"'name' in globals()".
+
+The Symbols that are created using this have default assumptions.
+If you want to place assumptions on symbols, you should create them
+using symbols() or var().
+
+Finally, this only works in the top level namespace. So, for
+example, if you define a function in isympy with an undefined
+Symbol, it will not work.
+.TP
+\*(T<\fB\-D\fR\*(T>, \*(T<\fB\-\-debug\fR\*(T>
+Enable debugging output. This is the same as setting the
+environment variable SYMPY_DEBUG to 'True'. The debug status is set
+in the variable SYMPY_DEBUG within isympy.
+.TP
+-- \fIPYTHONOPTIONS\fR
+These options will be passed on to \fIipython (1)\fR shell.
+Only supported when ipython is being used (standard python shell not supported).
+
+Two dashes (--) are required to separate \fIPYTHONOPTIONS\fR
+from the other isympy options.
+
+For example, to run iSymPy without startup banner and colors:
+
+isympy -q -c ipython -- --colors=NoColor
+.TP
+\*(T<\fB\-h\fR\*(T>, \*(T<\fB\-\-help\fR\*(T>
+Print help output and exit.
+.TP
+\*(T<\fB\-v\fR\*(T>, \*(T<\fB\-\-version\fR\*(T>
+Print isympy version information and exit.
+.SH FILES
+.TP
+\*(T<\fI${HOME}/.sympy\-history\fR\*(T>
+Saves the history of commands when using the python
+shell as backend.
+.SH BUGS
+The upstreams BTS can be found at \(lahttps://github.com/sympy/sympy/issues\(ra
+Please report all bugs that you find in there, this will help improve
+the overall quality of SymPy.
+.SH "SEE ALSO"
+\fBipython\fR(1), \fBpython\fR(1)

mnist_mlp/test.py

+import torch
+from gen_weight import *
+from main import Net
+from torchvision import datasets, transforms
+
+model_path = "./mnist_mlp_little_net.pt"
+
+# test script, not usefull for the lab work
+
+if __name__ == "__main__" :
+    model = Net()
+    model.load_state_dict(torch.load(model_path, weights_only=True))
+    max = find_max(model)
+
+    transform=transforms.Compose([
+        transforms.ToTensor(),
+        ])
+    dataset2 = datasets.MNIST('../data', train=False, transform=transform)
+    test_loader = torch.utils.data.DataLoader(dataset2, batch_size = 128)
+
+    gen_weight(model,bit_width = 32, int_width = 7, den_width = 25,dirname = "w_n_in")
+    gen_input(test_loader, num_samples = 10, bit_width = 32, int_width = 7, den_width = 25, dirname = "w_n_in")
+