thop计算模型复杂度(params,flops)

 thop安装

-pip install thop在线安装失败

-离线安装

github网址: 

pytorch-OpCounter:Count the MACs / FLOPs of your PyTorch model. - GitCode

python setup.py install

 测试:

from options import config as c
import os
os.environ["CUDA_VISIBLE_DEVICES"] = c.os_environ
import torch.nn
from modules.NET import Net
from utils.utils import load
from utils.yml import parse_yml, dict_to_nonedict
import numpy as np
from thop import profile
from modules.DCTGate_fast import DCT_transform

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# --------------------noise set------------------------------
# read noise_config
yml_path = c.noise_opt_yml_path
option_yml = parse_yml(yml_path)
# convert to NoneDict, which returns None for missing keys
noise_opt = dict_to_nonedict(option_yml)

# -------------------MODEL load---------------------------
model_path = c.load_model_path
net = Net(noise_opt, device).to(device)
load(model_path, net)

# -----------------MODEL input-----------------------
cover = torch.randn(1, 3, c.cropsize_val, c.cropsize_val).to(device)
secret = torch.Tensor(np.random.choice([-0.5, 0.5], (cover.shape[0], c.input_message_length))).to(device)

dct_trans = DCT_transform(image_size=c.cropsize_val, block_size=8).to(device)
cover_dct = dct_trans(cover)

# ------------cal: params, FLOPS-----------
flops, params = profile(net, (cover, secret, cover_dct))
print(f'\nflops: {flops}, params: {params}\n')
print('the flops is {}G, the params is {}M\n'.format(round(flops / (10 ** 9), 2), round(params / (10 ** 6), 2)))



评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值