from flops_counter import nn from flops_counter.tensorsize import TensorSize ###### # test on ReLU ###### relu = { 'layers': [ nn.ReLU() # same shape ], 'ins': [TensorSize([1, 64, 112, 112])], 'out_shape': [TensorSize([1, 64, 112, 112])], 'out_flops': [1605632] } test_on(relu) ###### # test on Sigmoid ###### sigmoid = { 'layers': [ nn.Sigmoid() # same shape ], 'ins': [TensorSize([1, 1, 56, 56])], 'out_shape': [TensorSize([1, 1, 56, 56])], 'out_flops': [9408] } test_on(sigmoid)
from _utils import test_on import sys sys.path.append('.') from flops_counter import nn from flops_counter.tensorsize import TensorSize ###### # test on Upsample ###### upsample = { 'layers': [ nn.Upsample(scale_factor=2, mode='bilinear') # same shape ], 'ins': [TensorSize([1, 1024, 20, 20])], 'out_shape': [TensorSize([1, 1024, 40, 40])], 'out_flops': [15974400] } test_on(upsample)
from _utils import test_on import sys sys.path.append('.') from flops_counter import nn from flops_counter.tensorsize import TensorSize ###### # test on Linear ###### linear = { 'layers': [ nn.Linear(4096, 8192) # same shape ], 'ins': [TensorSize([1, 4096])], 'out_shape': [TensorSize([1, 8192])], 'out_flops': [67108864] } test_on(linear)
from _utils import test_on import sys sys.path.append('.') from flops_counter import nn from flops_counter.tensorsize import TensorSize ###### # test on MaxPool2d ###### mxpool2d = { 'layers': [ nn.MaxPool2d(3, 2, 1) # same shape ], 'ins': [ TensorSize([1, 64, 112, 112]) ], 'out_shape': [ TensorSize([1, 64, 56, 56]) ], 'out_flops': [ 602112 ] } test_on(mxpool2d)
nn.Conv2d(64, 64, 3, 1, 1, bias=False) # same shape ], 'ins': [ TensorSize([1, 3, 224, 224]), TensorSize([1, 64, 56, 56]), TensorSize([1, 64, 56, 56]) ], 'out_shape': [ TensorSize([1, 64, 112, 112]), TensorSize([1, 64, 56, 56]), TensorSize([1, 64, 56, 56]) ], 'out_flops': [235225088, 25489408, 231010304] } test_on(conv) ###### # test on ConvTranspose2d ###### convtran = { 'layers': [ nn.ConvTranspose2d(512, 256, 4, 2, 1), # double the shape, except channels nn.ConvTranspose2d(1024, 256, 4, 4, 0) # quadro the shape, except channels ], 'ins': [TensorSize([1, 512, 28, 28]), TensorSize([1, 1024, 14, 14])], 'out_shape': [TensorSize([1, 256, 56, 56]), TensorSize([1, 256, 56, 56])],
eltadd = { 'layers': [ nn.EltAdd() # same shape ], 'ins': [ TensorSize([64, 112, 112]) ], 'out_shape': [ TensorSize([64, 112, 112]) ], 'out_flops': [ 802816 ] } test_on(eltadd) ###### # test on EltMul ###### eltmul = { 'layers': [ nn.EltMul() # same shape ], 'ins': [ TensorSize([1, 64, 112, 112]) ], 'out_shape': [ TensorSize([1, 64, 112, 112]) ], 'out_flops': [
sys.path.append('.') from flops_counter import nn from flops_counter.tensorsize import TensorSize ###### # test on BatchNorm2d ###### bn2d = { 'layers': [ nn.BatchNorm2d(64) # same shape ], 'ins': [TensorSize([1, 64, 112, 112])], 'out_shape': [TensorSize([1, 64, 112, 112])], 'out_flops': [4816896] } test_on(bn2d) ###### # test on L2Norm2d ###### l2norm2d = { 'layers': [ nn.L2Norm2d(256) # same shape ], 'ins': [TensorSize([1, 256, 56, 56])], 'out_shape': [TensorSize([1, 256, 56, 56])], 'out_flops': [2408448] } test_on(l2norm2d)