最近完善了一個(gè)Python庫(kù)，做到了使用一行命令便可以獲取到DNN網(wǎng)絡(luò)的大小情況。現(xiàn)在進(jìn)行發(fā)布，方便研究人員使用。

一、前言

為了獲取到神經(jīng)網(wǎng)絡(luò)模型的feature map以及weight的大小，我嘗試了很多方法，目前最方便的方法是用summary這個(gè)庫(kù)（https://github.com/Pinging-ZJU/Pytorch-Memory-Utils），然而在我們的測(cè)試下發(fā)現(xiàn)這個(gè)庫(kù)其實(shí)存在一些漏洞，使得對(duì)非線性網(wǎng)絡(luò)的測(cè)量不準(zhǔn)確（多測(cè)了很多層），于是我們對(duì)該庫(kù)進(jìn)行了修改，并添加了一些值方便用戶閱讀。

二、使用方法

1 直接在服務(wù)器上調(diào)用：pip3 install DNN_printer

2 在代碼中添加from DNN_printer import DNN_printer以及DNN_printer(net, (3, 32, 32),batch_size)
其中，第一個(gè)參數(shù)為net的定義，第二個(gè)(3,32,32)為輸出數(shù)據(jù)的維度，第三個(gè)數(shù)據(jù)為batch size的大小。

即：

from DNN_printer import DNN_printer

batch_size = 512
def train(epoch):
    print('\nEpoch: %d' % epoch)
    net.train()
    train_loss = 0
    correct = 0
    total = 0
    // put the code here and you can get the result
    DNN_printer(net, (3, 32, 32),batch_size)
    ...
    ...

之后就可以跑代碼了

三、結(jié)果

運(yùn)行代碼后得到如下打?。?/p>

image.png

即可以看到每一層的shape，大小（MB）以及weight的浮點(diǎn)數(shù)個(gè)數(shù)、大小（MB）信息。

GitHub庫(kù)如下：

https://github.com/Pinging-ZJU/DNN-Printer

如果好用歡迎大家給個(gè)star！謝謝哦，這里也要感謝Peiyi的合作