Save & load pre-trained model on Caffe2, Caffe2中儲存與讀取訓練過的模型

Caffe2官方教學中有介紹Loading Pre-Trained Models,以及在官方的Model Zoo中可以找到一些範例的model,而在model zoo中有提到在Caffe2中要load and inference pre-trained model你需要以下兩個檔案

predict_net.pb : 定義網路架構的protobuf file
init_net.pb : 紀錄該網路架構中weights的protobuf file
通常predict_net.pb會很小(KB級),而init_net.pb則是因為要紀錄weights所以大小大約是幾MB到上百MB不等,端看model的大小。

所以我們的目的是將訓練好的model輸出成init_net.pb 和 predict_net.pb,以下修改官方的mnist tutorial來測試。可以在training完成後透過SaveNet將想要trained model輸出成init_net.pb & predict_net.pb

def SaveNet(INIT_NET, PREDICT_NET, workspace, model):
init_net, predict_net = mobile_exporter.Export(
workspace, model.net, model.params
)

with open(PREDICT_NET, ‘wb’) as f:
f.write(model.net._net.SerializeToString())
with open(INIT_NET, ‘wb’) as f:
f.write(init_net.SerializeToString())

但由於不曉得為什麼,mnist tutorial中training完的test model accuracy只有0.1左右,所以先不管直接將train_model輸出成init_net.pb & predict_net.pb,

再透過LoadNet來讀出來測試。

def LoadNet(INIT_NET, PREDICT_NET, device_opts):
init_def = caffe2_pb2.NetDef()
with open(INIT_NET, 'r') as f:
init_def.ParseFromString(f.read())
init_def.device_option.CopyFrom(device_opts)
workspace.RunNetOnce(init_def.SerializeToString())

net_def = caffe2_pb2.NetDef()
with open(PREDICT_NET, 'r') as f:
net_def.ParseFromString(f.read())
net_def.device_option.CopyFrom(device_opts)
workspace.CreateNet(net_def.SerializeToString(), overwrite=True)

發生了下面的Runtime Error

RuntimeError: [enforce fail at operator.cc:26] blob != nullptr. op Iter: Encountered a non-existing input blob: ITER

這是因為init_net & predict_net應該只要inference的部份即可,不要將training過程會用到的ops加進去

所以要先解決為什麼test_model的accuracy這麼低的問題,

因為ModelHelper中的init_params=False沒有作用(official issue),所以若依照官方tutorial中在training後才call

# init test_model after training...
workspace.RunNetOnce(test_model.param_init_net)

則會使得param被重新init,導致最後test_model的testing accuracy只有0.1上下(因為weights都重新初始了,沒有繼承train_model訓練好的weights)。

workaround的解法是在training之前就先init test_model(總之就是train_model init的時候一起init test_model & deploy_model

# Initialize train_model, test_model and deploy_model
workspace.RunNetOnce(train_model.param_init_net)
workspace.RunNetOnce(test_model.param_init_net)
workspace.RunNetOnce(deploy_model.param_init_net)

這樣最後training完成後的testing accuracy就可以到0.95左右,

關於這個issue,在Caffe2 issue #824有PR了一筆commit是改在core.py中解了root cause,這邊先用worksaround解法來進行

在解決weights異常的問題後,可以在training結束將deploy_model輸出成init_net.pb & predict_net.pb,接著暫時隨機產生1x1x28x28, type=float32的array丟進去測試,

# generate a fake image with NCHW in float32
data = np.random.rand(1, 1, 28, 28).astype(np.float32)
# feed the blob and run the net
workspace.FeedBlob('data', data,device_option=device_opts)
workspace.RunNet('mnist_deploy', 1)

要注意的是,這邊的blob name ’data’ 和model name ‘mnist_deploy’要跟當初創建的時候相同(或是說要跟predict_net中的model name和input data blob name相同)。

You can find sample code here