pytorch计算flops, params，有两个input

茵茵同学

1794人浏览 · 2022-05-12 10:58:50

茵茵同学 · 2022-05-12 10:58:50 发布

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)

2048 AI社区

有“AI”的1024 = 2048，欢迎大家加入2048 AI社区

更多推荐

Kubernetes CoreDNS 解析超时排查（ndots 与 NodeLocal DNSCache）

CoreDNS 超时先用 ClusterIP 直连排除数据面，再读 Pod的ndots 与 search。默认 ndots:5 让短名查询放大数倍，AAAA 再乘一倍。指标上区分 NXDOMAIN 瀑布和真正的 SERVFAIL/p99。NodeLocal DNSCache 降 RTT、不消灭放大；根上仍是 FQDN 或谨慎的 dnsConfig。改 DNS 全局配置必须可回滚，禁止先停 Core