pytorch计算flops, params，有两个input

茵茵同学

已于 2022-05-12 10:59:36 修改

阅读量1.1k

点赞数 3

文章标签： pytorch 深度学习 python

于 2022-05-12 10:58:50 首次发布

本文链接：https://blog.csdn.net/pylyy/article/details/124726756

版权

该博客使用PyTorch库，通过CUDA设备进行模型运算加速。首先检查GPU是否可用，然后定义网络模型并将其转移到GPU上。接着，利用torchprofile模块对模型进行输入，并计算FLOPs和参数数量，以评估模型的计算复杂度和内存需求。

摘要由CSDN通过智能技术生成

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
net = D()  # 定义好的网络模型
if torch.cuda.is_available():
    net.cuda()
input1 = torch.randn(1, 3, 512, 512)#输入到模型时，模型要加一个维度x = x.unsqueeze(0)
input1 = input1.to(device)
#
input2 = torch.randn(1, 3, 512, 512)
input2 = input2.to(device)

flops, params = profile(net,inputs=(input1,input2))
print('flops: ', flops, 'params: ', params)