A python implementation for the PYNQ-BNN MNIST example

bnn_pynq.py

import numpy as np
from udmabuf import Udmabuf
from uio import Uio
import os
import struct
import math

# set path to the BNN-PYQN repo.
BNN_PYNQ_ROOT = '.'
PARAM         = 'bnn/params/mnist'
TEST_IMAGE    = 'tests/Test_image/3.image-idx3-ubyte'

### network parameters defined in "bnn/src/network/lfc-pynq/hw/config.h"
L0_SIMD=64
L0_PE=32
L0_WMEM=416
L0_TMEM=32
L0_MW=832
L0_MH=1024

L1_SIMD=32
L1_PE=64
L1_WMEM=512
L1_TMEM=16
L1_MW=1024
L1_MH=1024

L2_SIMD=64
L2_PE=32
L2_WMEM=512
L2_TMEM=32
L2_MW=1024
L2_MH=1024

L3_SIMD=8
L3_PE=16
L3_WMEM=128 # modified
L3_TMEM=1   # modified
L3_MW=1024
L3_MH=64

class BNN:
    def __init__(self):
        self.uio = Uio('uio0')
        self.mem = self.uio.memmap
        self.array   = np.frombuffer(self.mem, np.uint32, self.uio.length >> 2)

        self.udmabuf0 = Udmabuf('udmabuf0') # in
        self.udmabuf1 = Udmabuf('udmabuf1') # out

        n = 13
        self.buf_inp = self.udmabuf0.memmap(dtype=np.uint64, shape=(n))
        self.buf_out = self.udmabuf1.memmap(dtype=np.uint64, shape=(1))
        
        # num bytes & 1 : PC -> PL, 2 : PL -> PC
        self.udmabuf0.config_sync(8*n, 1)
        self.udmabuf1.config_sync(8*1, 2)
        
        self.reg_write64(0x10, self.udmabuf0.phys_addr)
        self.reg_write64(0x1c, self.udmabuf1.phys_addr)
        
    def reg_write(self,adr,data):
        self.array[adr >> 2] = data

    def reg_read(self,adr):
        return self.array[adr >> 2]

    def reg_write64(self,adr,data):
        ind = adr >> 2
        self.array[ind  ] =  data      & 0xffffffff
        self.array[ind+1] = (data>>32) & 0xffffffff
    
    def exec(self):
        self.reg_write(0x0, 0x1)
        while self.reg_read(0x0) & 0x2 == 0x1:
            pass
        return
        
    def doit(self, layer, mem, ind, data):
        self.reg_write(0x28, 0x1)    # doInit
        self.reg_write(0x30, layer)  # targetLayer
        self.reg_write(0x38, mem)    # targetMem
        self.reg_write(0x40, ind)    # targetInd

        self.reg_write64(0x48, data) # val

        self.exec()
        self.reg_write(0x28, 0x0)    # disable doInit
        
    def loadmem(self, layer, npe, lwmem, ltmem):
        for pe in range(0, npe):
            fn = self.param_dir + '/' + str(layer) + '-' + str(pe) + '-weights.bin'
            print(fn)
            with open(fn,'rb') as f:
                for line in range(0, lwmem):
                    c = f.read(8)
                    h = np.uint64(0x0)
                    if len(c) == 8:
                        h = struct.unpack('Q', c)[0]
                    else:
                        print("err")
                    self.doit(layer*2, pe, line, h)
                        
        for pe in range(0, npe):
            fn = self.param_dir + '/' + str(layer) + '-' + str(pe) + '-thres.bin'
            print(fn)
            with open(fn,'rb') as f:
                for line in range(0, ltmem):
                    c = f.read(8)
                    h = np.uint64(0x0)
                    if len(c) == 8:
                        h = struct.unpack('Q', c)[0]
                    else:
                        print("err")
                    self.doit(layer*2+1, pe, line, h)

    def load_parameters(self, dir):
        self.param_dir = dir
        print("setting weights and threasholds")
        self.loadmem(0, L0_PE, L0_WMEM, L0_TMEM)
        self.loadmem(1, L1_PE, L1_WMEM, L1_TMEM)
        self.loadmem(2, L2_PE, L2_WMEM, L2_TMEM)
        self.loadmem(3, L3_PE, L3_WMEM, L3_TMEM)

    def read_mnist(self, path):
        # https://github.com/lisa-lab/pylearn2/blob/master/pylearn2/utils/mnist_ubyte.py
        print(path)
        with open(path,'rb') as f:
            magic, number, rows, cols = struct.unpack('>iiii', f.read(16))
            print("header", magic, number, rows, cols)
            array = np.fromfile(f, dtype='uint8').reshape((number, rows, cols))
            
        # binarize as float32
        scale_min = -1.0
        scale_max = 1.0
        array_f = np.float32(array/255.0*(scale_max - scale_min) + scale_min)
        
        return array_f

    def read_mnist_labels(self, path):
        print(path)
        with open(path,'rb') as f:
            magic, number = struct.unpack('>ii', f.read(8))
            print("header", magic, number)
            array = np.fromfile(f, dtype='uint8').reshape((number))
        return array

    def paddedSize(self, x, padTo):
        # return in padded to a multiple of padTo
        if x % padTo == 0:
            return x
        else:
            return x + padTo - (x % padTo)

    def binarize(self, image_f):
        n64 = np.uint32(self.paddedSize(test_image_f[0].size,64)/64)

        image_b = np.zeros((n64), dtype=np.uint64)
        x = np.reshape(image_f, (image_f.size))
        for i in range(0, image_f.size):
            if x[i] > 0.0:
                image_b[int(i/64)] |= np.uint64(0x1 << (i % 64))

        return image_b
                
    def inference(self, test_image_f, count):
        test_image_b = bnn.binarize(test_image_f)

        # transefer image
        for i in range(0,13):
            self.buf_inp[i] = test_image_b[i]
        
        self.reg_write(0x34, 0x0)   # disable doInit 
        self.reg_write(0x54, count) # numReps

        self.udmabuf0.sync_for_device()

        self.exec()

        self.udmabuf1.sync_for_cpu()

        mask = (0x1 << 10) - 1
        result = int(self.buf_out[0]) & mask
        if result != 0:
            result = int(math.log2(float(result)))

        return result
        
bnn = BNN()
bnn.load_parameters(BNN_PYNQ_ROOT + '/' + PARAM)
test_image_f = bnn.read_mnist(BNN_PYNQ_ROOT + '/' + TEST_IMAGE)
result = bnn.inference(test_image_f[0], 1)
print("inference for the test image ", result)