NN-Tools 是一款方便易用的神经网络工具集，提供了面向 Caffe/PyTorch/TensorFlow 等主流神经网络框架的转换器/构建器/分析器。具有以下特性：
- 转换不同框架之间的模型。（Caffe，PyTorch，MXNet）
- 提供了方便的工具，可用于操作 caffe 模型和 ProtoTXT 文件（比如获取网络层的权重）。
- 分析模型，获取每一层的运算数。
模型转换
Pytorch to Caffe
这个工具支持 PyTorch 几乎所有的版本（从 0.2.0 到 0.4.0）。
- 支持的层: conv2d, linear, max_pool2d, avg_pool2d, dropout, relu, prelu, threshold(only value=0), softmax, batch_norm
- 支持的操作: torch.split, torch.max, torch.cat
- 支持的tensor操作: var.view, + (add), += (iadd), -(sub), -=(isub) * (mul) *= (imul)
示例：将 Alexnet 的 PyTorch 模型转换成 Caffe
1 2 3 4 5 6 7 8 9 10 11 12 13 14 | import sys sys.path.insert(0,'.') import torch from torch.autograd import Variable from torchvision.models.alexnet import alexnet import pytorch_to_caffe if __name__=='__main__': name='alexnet' net=alexnet(True) input=Variable(torch.ones([1,3,226,226])) pytorch_to_caffe.trans_net(net,input,name) pytorch_to_caffe.save_prototxt('{}.prototxt'.format(name)) pytorch_to_caffe.save_caffemodel('{}.caffemodel'.format(name)) |
示例：将 Inception3 的 PyTorch 模型转换成 Caffe
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 | import sys sys.path.insert(0,'.') import torch from torch.autograd import Variable from torchvision.models.inception import inception_v3 import pytorch_to_caffe if __name__=='__main__': name='inception_v3' net=inception_v3(True,transform_input=False) net.eval() input=Variable(torch.ones([1,3,299,299])) pytorch_to_caffe.trans_net(net,input,name) pytorch_to_caffe.save_prototxt('{}.prototxt'.format(name)) pytorch_to_caffe.save_caffemodel('{}.caffemodel'.format(name)) |
示例：将 VGG19 的 PyTorch 模型转换成 Caffe
1 2 3 4 5 6 7 8 9 10 11 12 13 14 | import sys sys.path.insert(0,'.') import torch from torch.autograd import Variable from torchvision.models.vgg import vgg19 import pytorch_to_caffe if __name__=='__main__': name='vgg19' net=vgg19(True) input=Variable(torch.ones([1,3,224,224])) pytorch_to_caffe.trans_net(net,input,name) pytorch_to_caffe.save_prototxt('{}.prototxt'.format(name)) pytorch_to_caffe.save_caffemodel('{}.caffemodel'.format(name)) |
工具地址：Github
本站微信群、QQ群(三群号 726282629)：
