brocolli

项目介绍

brocolli 是一个基于 torch fx 的工具，用于将 PyTorch 模型转换为其他格式，如 Caffe 和 ONNX。此外，该工具还包含 PyTorch 模型的量化功能。然而需要注意的是，目前 brocolli 项目已经停止维护，不再更新。

功能

PyTorch 模型转换

brocolli 提供了将 PyTorch 模型转换为 Caffe 和 ONNX 格式的功能，这对需要在不同框架之间迁移模型的用户非常有帮助。以下是两种转换方式的简要说明：

PyTorch 转 Caffe

要将 PyTorch 模型转换为 Caffe 格式，用户首先需要安装 brocolli-caffe。安装命令如下：

pip install brocolli-caffe

完成安装后，可以通过如下代码实现转换：

import torchvision.models as models
from brocolli.converter.pytorch_caffe_parser import PytorchCaffeParser

net = models.alexnet(pretrained=False)
x = torch.rand(1, 3, 224, 224)
pytorch_parser = PytorchCaffeParser(net, x)
pytorch_parser.convert()
pytorch_parser.save('alexnet')

运行脚本后，屏幕将显示 "accuracy test passed"，此时生成的 alexnet.caffemodel 和 alexnet.prototxt 文件会保存在当前文件夹中。

PyTorch 转 ONNX

将 PyTorch 模型转换为 ONNX 格式的过程与转换为 Caffe 类似。使用以下代码进行转换：

import torchvision.models as models
from brocolli.converter.pytorch_onnx_parser import PytorchOnnxParser

net = models.alexnet(pretrained=False)
x = torch.rand(1, 3, 224, 224)
pytorch_parser = PytorchOnnxParser(net, x)
pytorch_parser.convert()
pytorch_parser.save('alexnet.onnx')