pytorch计算flops, params，有两个input

原创已于 2022-05-12 10:59:36 修改 · 1.6k 阅读

7 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch #深度学习 #python

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

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

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

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)