TNN：由騰訊優圖實驗室打造，移動端高性能、輕量級推理框架，同時擁有跨平臺、高性能、模型壓縮、代碼裁剪等眾多突出優勢。TNN框架在原有Rapidnet、ncnn框架的基礎上進一步加強了移動端設備的支持以及性能優化，同時也借鑒了業界主流開源框架高性能和良好拓展性的優點。目前TNN已經在手Q、微視、P圖等應用中落地，歡迎大家參與協同共建，促進TNN推理框架進一步完善。
騰訊推理引擎TNN
本系列文章為對騰訊TNN的深度源碼級別解讀，希望通過對一個推理框架的完整描述，來增強讀者對于神經網絡設計、實現到優化的方方面面。
本節將對TNN中的優化策略管理執行過程做詳細的分析介紹。
在神經網絡優化管理器中主要定義了NetOptimizerManager和NetOptimizerRegister兩個類，分別實現管理執行和注冊的功能，接下來將對這兩個類的具體功能做詳細分析。
網絡優化策略管理器(NetOptimizerManager)
這里先來看看優化策略（網絡優化器net_optimizer）的定義。TNN中定義了多種神經網絡優化器（優化策略），比如fuse_conv_relu可以將conv層合并達到減少內存/顯存拷貝，減少計算和調用開銷的目的。NetOptimizer是所有這些優化器基類定義了優化器需要具備的基礎功能。具體的每種優化器優化方法后續文章依次說明。
class NetOptimizer {
public:
virtual ~NetOptimizer() {}
//定義策略
virtual std::string Strategy() = 0;
//確認設備是否支持優化
virtual bool SupportDevice(DeviceType device) = 0;
//基于網絡結構進行優化
virtual Status Optimize(NetStructure *structure, NetResource *resource) = 0;
};
優化策略方法說明
- Strategy方法：返回對應優化器的字符串名稱。
當前內部定義的一些名稱如下：
static const std::string kNetOptimizerFuseConvRelu = "net_optimizer_fuse_conv_relu";
static const std::string kNetOptimizerInsertReformat = "net_optimizer_Insert_reformat";
static const std::string kNetOptimizerRemoveLayers = "net_optimizer_remove_layers";
- SupportDevice方法：返回該優化器是否在設備上可用。
- Optimize方法:定義了具體的優化方法。
接下來看一下本文的重點NetOptimizerManager的實現。
NetOptimizerManager定義了優化管理器的執行方法。通過map和vector來存儲優化策略。其中map結構維護了strategy名稱到優化器實現的映射。vector結構維護了優先級信息和strategy名稱的映射。
//@brief net optimize: fuse relu and relu6 to convolution
class NetOptimizerManager {
public:
//執行優化
static Status Optimize(NetStructure *structure, NetResource *resource, DeviceType device);
//靜態方法：注冊優化策略并指定優先級
static void RegisterNetOptimizer(NetOptimizer *ptimizer, OptPriority prior);
private:
static std::map<std::string, std::shared_ptr<NetOptimizer>> &GetNetOptimizerMap();
static std::vector<std::pair<OptPriority, std::string>> &GetNetOptimizerSeq();
};
Optimize方法：
Optimize方法：定義了優化的執行邏輯框架（具體的優化由每個optimizer來執行）。
管理器需要依次執行每個優化器，這里是基于優先級來執行優化策略的。具體方法是通過對vector序列按優先級排序，然后遍歷map這樣可以按優先級從高到低依次執行優化。
支持的優先級，目前為P0到PLAST可以根據自己的需求進行擴展。
typedef enum {
// TOP
P0 = 0,
// MIDDLE
P1 = 1,
//
P2 = 2,
// LAST
PLAST = 1000
} OptPriority;
//基于優先級排序依次執行優化器。
Status NetOptimizerManager::Optimize(NetStructure *structure, NetResource *resource, DeviceType device) {
//獲取當前優化器的map結構
auto &optimizer_map = NetOptimizerManager::GetNetOptimizerMap();
//按照strategy字符串對優化器進行排序
std::sort(NetOptimizerManager::GetNetOptimizerSeq().begin(), NetOptimizerManager::GetNetOptimizerSeq().end());
//循環優化器
for (auto iter : NetOptimizerManager::GetNetOptimizerSeq()) {
//
auto optimizer = optimizer_map[iter.second];
//確認當前設備是否支持該優化
if (optimizer->SupportDevice(device)) {
//執行優化邏輯
auto status = optimizer->Optimize(structure, resource);
if (status != TNN_OK) { //執行出錯會導致優化中斷，返回對應狀態值
return status;
}
}
}
return TNN_OK;
}
RegisterNetOptimizer方法：
RegisterNetOptimizer方法定義了每個優化器optimizer的注冊過程，其實就是map和vector插入對應數據就OK了。
//參數：優化器和優先級
void NetOptimizerManager::RegisterNetOptimizer(NetOptimizer *optimizer, OptPriority prior) {
//優化器不為空
if (optimizer && optimizer->Strategy().length() > 0) {
auto &optimizer_map = NetOptimizerManager::GetNetOptimizerMap();
//往map和vector中將消息注冊進去
optimizer_map[optimizer->Strategy()] = std::shared_ptr<NetOptimizer>(optimizer);
NetOptimizerManager::GetNetOptimizerSeq().push_back(std::make_pair(prior, optimizer->Strategy()));
}
}
NetOptimizerRegister
在net_optimizer_manageer.h頭文件中，通過模板類將RegisterNetOptimizer方法封裝到NetOptimizerRegister類中，這樣就可以做到定義即注冊。
//通過模板指定具體的優化器。實現了定義即注冊的功能。
template <typename T>
class NetOptimizerRegister {
public:
explicit NetOptimizerRegister(OptPriority p) {
NetOptimizerManager::RegisterNetOptimizer(new T(), p);
}
};
