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)
测试模型参数量和flops
于 2024-02-01 15:13:50 首次发布