根據項目組要求，最近用Pytorch完成了U-net訓練BraTS2019數據的任務，輸出模型並在C++上重現部署，簡單說說我遇到的坑，可按目錄瀏覽：
libtorch中大大小小的坑：
- 模型轉換問題
步驟1：先在pytorch驗證你的模型(.pth)是否能夠重現結果，一般來講都可以成功復原的，這一步很重要，在C++上重現時最好有對照，以免思路混亂，這裏貼出我的測試一張圖片的部分代碼：
def load_model():
net = UNet2D(1, 5, 64 ).to(torch.device('cpu')) #cpu測試
state_dict = torch.load(saved_model_path,map_location='cpu')
net.load_state_dict(state_dict,strict=True) #記得加True
return net
if __name__ == "__main__":
np.set_printoptions(threshold=np.inf) #爲了看完整的結果
net = load_model()
img = sitk.ReadImage('BraTS19_2013_2_1_flair.nii.gz')#需要安裝SimpleITK這個包來讀nii
nda = sitk.GetArrayFromImage(img)
test_data = np.asarray(nda[110])
test_data = norm_vol(test_data)#這裏的歸一化自己寫的，按照自己需求來確定是否要做歸一化
#轉成tensor
test_data = torch.from_numpy(test_data)
test_data = torch.tensor(test_data,dtype=torch.float32)
#按照網絡輸入需求拓展維度
test_data=torch.unsqueeze(test_data,0)
test_data=torch.unsqueeze(test_data,1)
#預測階段
with torch.no_grad():
net.eval()
predict = net(test_data)
predict = F.softmax(predict,dim=1)
predict = torch.max(predict,dim=1)[1]
predict = predict.squeeze().long().data
io.imwrite('result.jpg', predict)#imageio工具包
步驟2：將pth模型轉至pt文件，後續用於C++預測，這個過程官方的例子裏也有，這是我按照自己需求來改的代碼：
import torch
import torch.nn as nn
from unet2d import UNet2D
saved_model_path = 'best.pth'
net = UNet2D(1, 5, 64 ).to(torch.device('cpu'))
state_dict = torch.load(saved_model_path,map_location='cpu')
net.load_state_dict(state_dict,strict=True)#同樣記得TRUE
net.eval()#重要！
example = torch.rand(1, 1, 240, 240).float()
traced = torch.jit.trace(net, example)
traced.save('best.pt')
其他：可以按壓縮包的打開方式來打開pt文件來查看模型追蹤是否正確：
其中這個文件可以查看整體的追蹤結果。
- tensor是nan的值
libtorch的安裝以及cmake運行在我之前的博客裏有提到，可以翻一下對號入座來安裝，網上教程也很多，這裏不加贅述，我pytorch和libtorch都是1.4.0。
部分C++部署代碼如下：
int main() {
auto tensor1 = torch::empty(1 * 1 * 240 * 240);
float* data1 = tensor1.data<float>();
for (int i = 0; i < 1; i++)
{
for (int j = 0; j < 1; j++)
{
for (int x = 0; x < 240; x++)
{
for (int y = 0; y < 240; y++)
{
*data1++ = (itk[110][x][y]-min*1.0)/(max.0-1.0);
}
}
}
}
auto t = tensor1.resize_({ 1,1,240,240 });
t=t.div(255);
torch::jit::script::Module module = torch::jit::load(model_path,torch::kCPU); //load model
// init model
module.eval();
cout << "model input is ok\n";
vector<torch::jit::IValue> inputs; //def an input
inputs.emplace_back(t.toType(torch::kFloat32));
float start = getTickCount();
auto result = module.forward( inputs ).toTensor(); //前向傳播獲取結果 = net(image)
float end = getTickCount();
float last = end - start;
cout << "time consume: " << (last / getTickFrequency()) << endl;
//rescalling input element into range(0,1) and sum to 1; output size is same to input
auto prob = result.softmax(1);//torch::nn::functional::softmax(result, 1);
auto prediction = prob.max(1);//tuple類型
inputs.pop_back();
std::tuple_element<1, decltype(prediction)>::type cnt = std::get<1>(prediction);
std::cout << "cnt = " << cnt.sizes() << std::endl;// 1 240 240
cnt=cnt.squeeze().data;//這句代碼有問題，最好按照需求來確認要不要寫
cout << "result sizes:" << cnt.sizes() << endl; //240 240
Mat m(cnt.size(0), cnt.size(1), CV_32FC1, cnt.data());
imwrite("res.jpg", m);
return 0;
}
從上述代碼可以看出流程：
讀入數據->載入模型->將數據整理成tensor(這裏我沒用torch::from_blob)->預測->結果整理->輸出
大體流程與上面的pytorch流程一致。
但在前向傳播這一步的時候，結果輸出nan值：
很明顯是不正確的，問題是出在我加了一句t=t.div(255);，導致我數據類型出錯，問題發生點可以看我在pytorch論壇裏的提問：
Pytorch Forums–Having problems in segmenting image when using libtorch
- 輸出結果只有部分的一塊
結果輸出有很多情況，就我而言，遇到的問題是輸出全黑的圖片和只有部分分割結果的圖片：
這種情況最好檢查輸入tensor的值，是否與pytorch上輸出的一致
這種情況最好確認是否需要加==cnt.squeeze.data();==這一句代碼，我刪除以後就沒問題了。
因爲我只訓練了50個epoch，所以結果很差，將就着先用。
先寫到這，想到再補充。