TNN iOS非图像模型入门

注:本文同步发布于微信公众号:stringwu的互联网杂谈TNN iOS 非图像模型入门指南

1 背景

TNN是腾讯优图实验室开源的高性能、轻量级神经网络推理框架TNN,github上也有比较详细的例子来说明如何在端上运行图像类的模型,但demo 更多是图像类相关的示例,而且里面做了一层层的封装,很难让一个初学者直接上手一步步构建出可推理的结果,
本文主要从初学者的角度出发,按照TNNAPI文档一步步构建出非图像模型的入门文档。(本文不再详述如何编译和集成TNN工程,有需要的同学可直接参考Demo文档);

2 构建

TNN的推理流程主要包括模型的解析,网络构建,输入设定,输出获取

2.1 模型解析

模型文件包括两个部分:

  • *.tnnmodel
  • *.tnnproto

模型解析的步骤包括:

  • 获取模型文件的路径
  • 解析文件内容
  • 初始化模型
	// 获取模型文件 (提前把对应的模型文件集成到工程中)
    auto model_path = [[NSBundle mainBundle] pathForResource:@"model/recommend/recommend.tnnmodel"
                                                      ofType:nil];
    auto proto_path = [[NSBundle mainBundle] pathForResource:@"model/recommend/recommend.tnnproto" ofType:nil];

    //解析文件内容
    string proto_content = [NSString stringWithContentsOfFile:proto_path encoding:NSUTF8StringEncoding error:nil].UTF8String;
    
    string model_content = [data_mode length] > 0 ? string((const char *)[data_mode bytes], [data_mode length]) : "";

    //模型的配置
    TNN_NS::ModelConfig model_config;
    model_config.model_type = TNN_NS::MODEL_TYPE_TNN;  // 指定模型的类型为TNN
    model_config.params.push_back(proto_content);
    model_config.params.push_back(model_content);

    auto tnn = std::make_shared<TNN_NS::TNN>(); //实例化TNN 的实例
    //初始化模型
    Status ret = tnn->Init(model_config);
    //结果为TNN_OK时才为模型初始化成功
    if (ret != TNN_OK) {
        return;
    }

2.2 网络构建

网络的构建需要配置TNN_NS::NetworkConfig,这个配置需要指定device_typelibrary_path,在iOS中的device_type正常是使用TNN_NS::DEVICE_ARMTNN_NS::DEVICE_METAL就可以了,但笔者在实际尝试时,发现device_type指定这两个类型都没有办法正常跑通,后与TNN相关同学咨询请教后,使用了TNN_NS::DEVICE_NAIVE才正常跑通,具体的原因TNN的同学还在帮忙定位中。(如果发现数据正常时,流程没有办法跑通的话,可以多换几个device_type看看)

  auto library_path = [[NSBundle mainBundle] pathForResource:@"tnn.metallib" ofType:nil];
  //shape数据
  auto shape_path = [[NSBundle mainBundle] pathForResource:@"model/recommend/tnn_input.json" ofType:nil];
  //将shape数据转成NSDictionary
  NSData *data_shape = [NSData dataWithContentsOfFile:shape_path];
  NSDictionary * dictionary_shape = [NSJSONSerialization JSONObjectWithData:data_shape options:NSJSONReadingMutableLeaves error:nil];
   //构造出TNN可识别的shape数据
  InputShapesMap sdkInputShape = {};
   for (NSString *key in dictionary_shape) {
        NSDictionary*tmpValue = dictionary_shape[key];
        NSLog(@"InputShapesMap: println key %@****",key);
        NSInteger int1 = [tmpValue[@"dim1"]intValue];
        NSInteger int2 = [tmpValue[@"dim2"]intValue];
        NSLog(@"InputShapesMap222: println key %ld,%ld****",int1,int2);
        TNN_NS::DimsVector nc = {(int)int1,(int)int2};
        sdkInputShape.insert(std::pair<std::string, TNN_NS::DimsVector>(std::string(key.UTF8String),nc));
    }


   // 构造出net_config
    TNN_NS::NetworkConfig net_config;
    net_config.device_type = TNN_NS::DEVICE_NAIVE; // 指定device_type,如果跑不成功,可以多换几个type试试
    net_config.library_path = {library_path.UTF8String};

   TNN_NS::Status error;
   //构造出TNN网络对象
   auto net_instace = tnn->CreateInst(net_config, error,sdkOptions->input_shapes);
     //结果为TNN_OK时才为网络构建成功
    if (error != TNN_OK) {
        return;
    }


2.3 输入设定

输入设定主要是通过TNN的方法 SetInputMat 来完成的;

//从文件获取模型的输入
auto mock_input_path = [[NSBundle mainBundle] pathForResource:@"model/recommend/mock_input.json" ofType:nil];
//将数据转换成NSDictionary
NSData *data_mock = [NSData dataWithContentsOfFile:mock_input_path];
NSDictionary * dictionary_mock = [NSJSONSerialization JSONObjectWithData:data_mock options:NSJSONReadingMutableLeaves error:nil];

//将数据转换成TNN的输入格式
std:map<std::string,std::vector<float>> inputDatas = {};
for (NSString *key in dictionary_mock) {
    NSArray *valueArray = dictionary_mock[key];
    __block std::vector<float> valueVector ={};
    valueVector.reserve([valueArray count]);
    [valueArray enumerateObjectsUsingBlock:^( id _Nonnull obj, NSUInteger idx, BOOL * _Nonnull stop) {
          valueVector.push_back([obj floatValue]);
    }];
    inputDatas[[key UTF8String]] = valueVector;
  }

BlobMap blob_map;
net_instace ->GetAllInputBlobs(blob_map);
//遍历所有的输入key
for (const auto&item : blob_map) {
    //获取对应的key
    std::string name = item.first;
    //获取key对应的输入数据
    std::vector<float> tmpItem = inputDatas.at(name);
    //获取key对应的shape数据
    DimsVector shape = sdkOptions->input_shapes[name];
    //构造出input_mat
    auto input_mat = std::make_shared<TNN_NS::Mat>(net_config.device_type,TNN_NS::NCHW_FLOAT,shape,tmpItem.data());

    MatConvertParam input_convert_params = TNN_NS::MatConvertParam();
    //把输入数据通过SetInputMat方法给到TNN引擎
    auto status = net_instace->SetInputMat(input_mat,input_convert_params,name);
    //结果为TNN_OK时才为设置输入成功
    if (status != TNN_OK) {
        NSLog(@"setInputmat error ");
        return;
     }

    }

//执行推理
auto forwardStatus = net_instace->Forward();
if (forwardStatus != TNN_OK) {
    NSLog(@"forwardStatus error ");
    return;
}

2.4 输出获取

获取推理的结果,是通过TNNGetAllOutputBlobs接口来完成


BlobMap out_blob_map;
//获取所有的输出key
net_instace ->GetAllOutputBlobs(out_blob_map);
//所有的结果都输出放在 mat_map里面
std::map<std::string, std::shared_ptr<TNN_NS::Mat> > mat_map = {};
//遍历所有输出key
for (const auto&item: out_blob_map) {
    auto name = item.first;
    MatConvertParam output_convert_params = TNN_NS::MatConvertParam();
    std::shared_ptr<TNN_NS::Mat> output_mat = nullptr;
    auto status = net_instace->GetOutputMat(output_mat,output_convert_params,name,net_config.device_type);
    if (status != TNN_OK) {
        NSLog(@"outPutError error ");
        return;
    }
    mat_map[name] = output_mat;
    NSLog(@"outPutname %@",[NSString stringWithFormat:@"%s", name.c_str()]);
}

在得到输出结果的map后mat_map后,就可以根据业务的情况进行结果的解析,本文使用的模型最终的结果是一个float值

//推荐模型的推理结果是存在在pred字段里面;
auto scores = mat_map["pred"];
auto dims = scores ->GetDims();
if (dims.size() <= 0) {
    NSLog(@"scores dims  error ");
    return;
}

//获取当前的推理结果
float *score_data = static_cast<float*>(scores->GetData());
NSLog(@"scores is %f ",score_data[0]);