测试模型参数量和flops

onlyzzr

已于 2024-04-24 16:41:17 修改

阅读量154

点赞数 3

分类专栏：深度学习文章标签：深度学习 pytorch 人工智能

于 2024-02-01 15:13:50 首次发布

本文链接：https://blog.csdn.net/weixin_45850570/article/details/135974733

版权

深度学习专栏收录该内容

4 篇文章 0 订阅

订阅专栏

from models_guided import Generator_F2S, Generator_S2F
from thop import profile
import torch
genA2B = Generator_S2F(3, 3,n_residual_blocks=9)
input = torch.randn(1, 3, 256, 256)
netG_A2B_FLOPs, netG_A2B_params = profile(genA2B, inputs=(input, ))
print(f"Total FLOPs: {(netG_A2B_FLOPs) / 1e9} GFLOPs")
print(f"Total Params: {(netG_A2B_params) / 1e6} MParams")
def print_network(net):
    num_params = 0
    for param in net.parameters():
        num_params += param.numel()
    print('Total number of parameters: %d' % num_params)