conv_size.py

import numpy as np


class Vector:
    def __init__(self, w, h, d, n=None):
        self.w = w
        self.h = h
        self.d = d
        self.n = n
        self.half_w = w // 2
        self.half_h = h // 2
        self.product = w * h * d

    def __str__(self):
        return 'w: {:d}, h: {:d}, d: {:d}, n: {:d}'.format(self.w, self.h, self.d, self.n) if self.n is not None else \
               'w: {:d}, h: {:d}, d: {:d}'.format(self.w, self.h, self.d)


def main():
    inputs = Vector(224, 224, 3)
    filters = Vector(7, 7, 3, 64)
    stride = 2
    padding = 3
    dilation = 0
    use_bias = False

    assert inputs.d == filters.d

    h = (inputs.h + 2 * padding - 2 * filters.half_h * (dilation + 1)) // stride
    w = (inputs.w + 2 * padding - 2 * filters.half_w * (dilation + 1)) // stride
    outputs = Vector(w, h, filters.n)

    fma = outputs.product * filters.product
    add = 0
    stores = outputs.product
    loads = filters.product * filters.n  # Load filters
    loads += inputs.product  # Load inputs TODO dilation

    if use_bias:
        add += outputs.product
        loads += filters.n

    print('Outputs', outputs)
    print('Inputs', inputs)
    print('Unique Loads', loads)
    print('Unique Stores', stores)
    print('FMAs', fma)
    print('ADDs', add)
    print('Flops', fma * 2 + add)
    print('FMAs per load', fma / loads)
    print('FMAs per store', fma / stores)


if __name__ == '__main__':
    main()