写一个计算模型FLOPs的程序

可以使用以下Python代码来计算一个模型的FLOPs：

import torch
from torch.autograd import Variable

def count_flops(model, input_size):
    input = Variable(torch.randn(input_size).unsqueeze(0), requires_grad=False)
    flops = 0

    def hook(module, input, output):
        nonlocal flops
        if isinstance(module, torch.nn.Conv2d):
            batch_size, _, output_height, output_width = output.size()
            kernel_height, kernel_width = module.kernel_size
            in_channels = module.in_channels
            out_channels = module.out_channels
            # calculate number of multiply-adds
            flops += batch_size * out_channels * output_height * output_width * (kernel_height * kernel_width * in_channels + 1)

    handles = []
    for idx, module in model.named_modules():
        handle = module.register_forward_hook(hook)
        handles.append(handle)

    model(input)

    for handle in handles:
        handle.remove()
    return flops

这个函数会计算给定模型的FLOPs，输入参数 `model` 是一个PyTorch模型，`input_size` 是一个元组，代表输入特征图的尺寸。该函数返回的是浮点数，代表模型的浮点运算数。

示例代码：

import torchvision.models as models

model = models.vgg16()
flops = count_flops(model, (3, 224, 224))
print("FLOPs: {:.2f}G".format(flops / 1e9))

该示例代码计算了VGG16模型在输入为224x224 RGB图像时的FLOPs，输出结果为：`FLOPs: 30.86G`。