快捷導(dǎo)航

Python中Thop庫(kù)的基本用例和參數(shù)說(shuō)明

更新時(shí)間：2024年05月21日 11:26:27 作者：肆十二

這篇文章主要給大家介紹了關(guān)于Python中Thop庫(kù)的基本用例和參數(shù)說(shuō)明的相關(guān)資料,THOP是PyTorch非常實(shí)用的一個(gè)第三方庫(kù),可以統(tǒng)計(jì)模型的 FLOPs 和參數(shù)量,文中通過(guò)代碼介紹的非常詳細(xì),需要的朋友可以參考下

Thop介紹

THOP（Torch-OpCounter）是一個(gè)用于計(jì)算PyTorch模型操作數(shù)和計(jì)算量的工具。通過(guò)這個(gè)庫(kù)，開(kāi)發(fā)者可以更好地理解和評(píng)估模型的復(fù)雜度，這對(duì)于模型優(yōu)化和性能調(diào)優(yōu)是非常有幫助的。

具體來(lái)說(shuō)，使用THOP，可以對(duì)定義的PyTorch模型進(jìn)行分析，統(tǒng)計(jì)出模型所包含的各種操作的數(shù)量以及模型的計(jì)算量。這對(duì)于理解模型的運(yùn)行效率、找出可能的性能瓶頸以及優(yōu)化模型結(jié)構(gòu)都是非常重要的。

在Python中使用THOP庫(kù)的基本步驟包括：首先安裝THOP庫(kù)，這通常可以通過(guò)pip命令進(jìn)行安裝；然后在Python腳本中導(dǎo)入THOP庫(kù)；最后定義需要分析的PyTorch模型，并使用THOP進(jìn)行模型的操作數(shù)和計(jì)算量的統(tǒng)計(jì)。

需要注意的是，THOP庫(kù)的具體使用方法和支持的功能可能會(huì)隨著版本的更新而發(fā)生變化，因此建議查閱最新的官方文檔以獲取最準(zhǔn)確的信息。

參數(shù)

model（PyTorch模型）：需要計(jì)算操作數(shù)和計(jì)算量的PyTorch模型。
input_size（輸入大?。耗Ｐ偷妮斎霃埩康拇笮?。這通常是一個(gè)元組，指定了輸入張量的形狀。
custom_ops（自定義操作）：一個(gè)字典，用于指定自定義操作的計(jì)算量。這可以用于覆蓋THOP庫(kù)中默認(rèn)的操作計(jì)算方式。
ignore_ops（忽略操作）：一個(gè)列表，指定了在計(jì)算過(guò)程中需要忽略的操作類型。這對(duì)于排除某些不重要的操作或特定層的計(jì)算量非常有用。

基本用例

安裝THOP庫(kù)：首先，你需要安裝THOP庫(kù)。通常，你可以使用pip命令進(jìn)行安裝：

pip install thop

導(dǎo)入必要的庫(kù)：在你的Python腳本中，導(dǎo)入必要的庫(kù)和模塊：

import torch  
import thop

定義PyTorch模型：定義一個(gè)PyTorch模型，該模型將用于計(jì)算操作數(shù)和計(jì)算量。
計(jì)算操作數(shù)和計(jì)算量：使用THOP庫(kù)的profile函數(shù)來(lái)計(jì)算模型的操作數(shù)和計(jì)算量。你需要提供模型和輸入大小作為參數(shù)。例如：

model = ...  # 定義你的PyTorch模型  
input_size = (1, 3, 224, 224)  # 輸入大小，這里假設(shè)是一個(gè)批量的3通道224x224圖像  
  
flops, params = thop.profile(model, inputs=(torch.randn(input_size),))  
print(f"FLOPs: {flops / 1e9} G")  # 打印計(jì)算量（以十億次浮點(diǎn)運(yùn)算為單位）  
print(f"Params: {params / 1e6} M")  # 打印參數(shù)量（以百萬(wàn)為單位）

在這個(gè)例子中，thop.profile函數(shù)返回了兩個(gè)值：flops表示模型的計(jì)算量（浮點(diǎn)運(yùn)算次數(shù)），params表示模型的參數(shù)量。這兩個(gè)值通常用于評(píng)估模型的復(fù)雜度和性能。注意，在計(jì)算FLOPs時(shí)，我們通常將其除以1e9來(lái)將其轉(zhuǎn)換為以十億次浮點(diǎn)運(yùn)算為單位（GFLOPs），同樣地，參數(shù)量通常除以1e6來(lái)將其轉(zhuǎn)換為以百萬(wàn)為單位（MParams）。

附：thop的安裝問(wèn)題

在python中安裝過(guò)第三方庫(kù)的小伙伴都知道，它的安裝命令是：pip install XXX。然而針對(duì)Thop的安裝卻有一個(gè)大坑。

1）直接使用常用命令

pip install thop

它會(huì)默認(rèn)安裝0.0.31.post2005241907的版本。由于我使用的pytorch環(huán)境為1.0.0，不支持： nn.SyncBatchNorm（多GPU執(zhí)行所需的），而該版本下的THop【THOP:0.031】沒(méi)有對(duì)pytorch:1.0.0修復(fù)此bug，從而導(dǎo)致運(yùn)行出錯(cuò)。