#include<torch/torch.h> #include<iostream> #include<vector> #include<cassert> #include<stdlib.h> #include<unordered_map> #include<fstream> class double_conv:public torch::nn::Module { public: torch::nn::Conv2d conv1,conv2; torch::nn::BatchNorm bn1,bn2; int in_ch,out_ch; public: double_conv(int in_ch,int out_ch):in_ch(in_ch),out_ch(out_ch),conv1(torch::nn::Conv2dOptions(in_ch,out_ch,3).padding(1)),bn1(out_ch), conv2(torch::nn::Conv2dOptions(out_ch,out_ch,3).padding(1)),bn2(out_ch) { register_module("conv1",conv1); register_module("conv2",conv2); register_module("bn1",bn1); register_module("bn2",bn2); } torch::Tensor forward(torch::Tensor x) { x = conv1->forward(x); x = bn1->forward(x); x = torch::relu(x); x = conv2->forward(x); x = bn2->forward(x); x = torch::relu(x); return x; } }; class inconv:public torch::nn::Module { public: int in_ch,out_ch; public: inconv(int in_ch,int out_ch):in_ch(in_ch),out_ch(out_ch){} torch::Tensor forward(torch::Tensor x) { double_conv dc(in_ch,out_ch); x = dc.forward(x); return x; } }; class down:public torch::nn::Module { public: int in_ch,out_ch; public: down(int in_ch,int out_ch):in_ch(in_ch),out_ch(out_ch){} torch::Tensor forward(torch::Tensor x) { x = torch::max_pool2d(x,2); double_conv dc(in_ch,out_ch); x = dc.forward(x); return x; } }; class up:public torch::nn::Module { public: int in_ch,out_ch; torch::nn::Conv2d upconv; torch::nn::Conv2d conv1,conv2; torch::nn::BatchNorm bn1,bn2; torch::Tensor x; public: up(int in_ch,int out_ch):in_ch(in_ch),out_ch(out_ch),upconv(torch::nn::Conv2dOptions(in_ch,out_ch,4).padding(1).stride(2).transposed(new bool(true))), conv1(torch::nn::Conv2dOptions(out_ch,out_ch,3).padding(1)),bn1(out_ch),conv2(torch::nn::Conv2dOptions(out_ch,out_ch,3).padding(1)),bn2(out_ch) { register_module("upconv",upconv); register_module("conv1",conv2); register_module("conv2",conv2); register_module("bn1",bn1); register_module("bn2",bn2); } torch::Tensor forward(torch::Tensor x1,torch::Tensor x2) { x = upconv->forward(x1); x = torch::cat({x,x2},1); double_conv dc(x.size(1),out_ch); x = dc.forward(x); //x = conv1->forward(x); //x = bn1->forward(x); //x = torch::relu(x); //x = conv2->forward(x); //x = bn2->forward(x); //x = torch::relu(x); return x; } }; class outconv:public torch::nn::Module { public: int in_ch,out_ch; torch::nn::Conv2d conv; public: outconv(int in_ch,int out_ch):in_ch(in_ch),out_ch(out_ch),conv(torch::nn::Conv2dOptions(in_ch,out_ch,1).padding(0)) { register_module("conv",conv); } torch::Tensor forward(torch::Tensor x) { return conv->forward(x); } }; class unet:public torch::nn::Module { public: int n_ch,n_class; inconv *iconv= new inconv(n_ch,64); down *down1= new down(64,256); down *down2= new down(256,512); down *down3= new down(512,512); down *down4= new down(512,512); up *up1= new up(512,256); up *up2= new up(256,128); up *up3= new up(128,64); up *up4= new up(64,64); outconv *oconv= new outconv(64,n_class); torch::Tensor x1,x2,x3,x4,x5; public: unet(int n_ch,int n_class):n_ch(n_ch),n_class(n_class){} torch::Tensor forward(torch::Tensor x) { x1 = iconv->forward(x); x2 = down1->forward(x1); x3 = down2->forward(x2); x4 = down3->forward(x3); x5 = down4->forward(x4); x = up1->forward(x5,x4); x = up2->forward(x,x3); x = up3->forward(x,x2); x = up4->forward(x,x1); x = oconv->forward(x); return x; } }; std::vector<float> Tokenize(const std::string& str,const std::string& delimiters) { std::vector<float> tokens; std::string::size_type lastPos = str.find_first_not_of(delimiters, 0); std::string::size_type pos = str.find_first_of(delimiters, lastPos); while (std::string::npos != pos || std::string::npos != lastPos) { tokens.push_back(std::atof(str.substr(lastPos, pos - lastPos).c_str())); lastPos = str.find_first_not_of(delimiters, pos); pos = str.find_first_of(delimiters, lastPos); } return tokens; } std::vector<std::vector<float>> readTxt(std::string file) { std::ifstream infile; infile.open(file.data()); assert(infile.is_open()); std::string s; std::vector<float> vec; std::vector<std::vector<float>> res; while(getline(infile,s)) { std::string tt= static_cast<std::string>(s); vec = Tokenize(tt, " "); res.push_back(vec); } infile.close(); std::cout<<"gdood"<<std::endl; return res; } torch::Tensor float2TensorLabel() { static float tt[2478][3125]={0}; //memset(tt,0,sizeof(tt)); std::vector<std::vector<float>> vec = readTxt("/Users/yanlang/unet/mx-unet/U-Net/LabelData.txt"); int ch = vec.size(); int len = vec[0].size(); for(int i=0;i<ch;i++) { for(int j=0;j<len;j++) { tt[i][j]=vec[i][j]; } } torch::Tensor tmask = torch::CPU(torch::kFloat).tensorFromBlob(tt,{2478,3125}); return tmask; } torch::Tensor float2TensorData() { static float tt[7][2478*3125] = {0}; std::vector<std::vector<float>> vec = readTxt("/Users/yanlang/unet/mx-unet/U-Net/ImageData.txt"); int ch = vec.size(); int len = vec[0].size(); for(int i=0;i<ch;i++) { for(int j=0;j<len;j++) { tt[i][j]=vec[i][j]; } } torch::Tensor tdata = torch::CPU(torch::kFloat).tensorFromBlob(tt,{7,2478,3125}); return tdata; } int imgH=256; int imgW=256; torch::Tensor RandData(torch::Tensor data,int hight,int width) { //torch::Tensor datat = torch::squeeze(data); torch::Tensor tmp = torch::zeros({7,imgH,imgW}); for(int i=hight;i<hight+imgH;i++) { for(int j=width;j<width+imgW;j++) { for(int k=0;k<7;k++) { tmp[k][i-hight][j-width]=data[k][i][j]; } } } return tmp; } torch::Tensor RandMask(torch::Tensor label, int hight,int width) { torch::Tensor tmp = torch::zeros({imgH,imgW}); for(int i=hight;i<hight+imgH;i++) { for(int j=width;j<width+imgW;j++) { tmp[i-hight][j-width]=label[i][j]; } } return tmp; } std::vector<torch::Tensor> DataLoader(torch::Tensor data,torch::Tensor label,int batch_size) { int imghight = data.size(1); int imgwidth = data.size(2); int randhight,randwidth; torch::Tensor resdata = torch::zeros({batch_size,7,imgH,imgW}); torch::Tensor reslabel = torch::zeros({batch_size,imgH,imgW}); for(int i=0;i<batch_size;i++) { randhight = rand()%(imghight-imgH-1); randwidth = rand()%(imgwidth-imgW-1); resdata[i] = RandData(data,randhight,randwidth); reslabel[i] = RandMask(label,randhight,randwidth); } return {resdata,reslabel}; } torch::autograd::Variable Get_predData(torch::autograd::Variable data) { //torch::autograd::Variable datat = torch::unsqueeze(data,0); torch::autograd::Variable tmp = torch::zeros({7,imgH,imgW}); for(int i=500;i<756;i++) { for(int j=500;j<756;j++) { for(int k=0;k<7;k++) { tmp[k][i-500][j-500]=data[k][i][j]; } } } return torch::unsqueeze(tmp,0); } void write2Txt(torch::autograd::Variable data) { std::ofstream fout("tresult.txt"); for(int i=0;i<data.size(0);i++) { for(int j=0;j<data.size(1);j++) { fout<<data[i][j]<<std::endl; } } fout.close(); } void saveModel(std::vector<torch::Tensor> weights,std::vector<std::string> key) { std::ofstream fout("unet.txt"); //std::unordered_map<std::string,torch::Tensor> mp; for(int i=0;i<weights.size();i++) { fout<<key[i]<<std::endl; fout<<weights[i]<<std::endl; } fout.close(); } void trainConvNet(unet model) { torch::optim::SGD optimizer(model.parameters(),/*lr=*/0.01); torch::Tensor pred; std::cout<<"load data ......"<<std::endl; torch::autograd::Variable data = torch::autograd::make_variable(float2TensorData()); torch::autograd::Variable label = torch::autograd::make_variable(float2TensorLabel()); std::cout<<"done!!"<<std::endl; torch::Tensor train_data,train_label; std::vector<torch::Tensor> vecdata; for(int epoch=0;epoch<20;epoch++) { vecdata = DataLoader(data,label,2); std::cout<<"vecdata after done!!"<<std::endl; train_data = vecdata[0]; std::cout<<"train_data after done"<<std::endl; train_label = vecdata[1]; std::cout<<train_label.size(0)<<std::endl; std::cout<<"train_label after done"<<std::endl; pred = model.forward(train_data); auto loss = torch::nll_loss2d(pred,torch::_cast_Long(train_label));//torch::_cast_Long() std::cout<<"the loss is"<<loss<<std::endl; optimizer.zero_grad(); loss.backward(); optimizer.step(); } std::vector<torch::Tensor> vecValue; std::vector<std::string> vecKey; torch::nn::ParameterCursor tt = model.parameters(); for(auto it=tt.begin();it!=tt.end();it++) { vecValue.push_back((*it).value); vecKey.push_back((*it).key); } saveModel(vecValue,vecKey); torch::autograd::Variable predData = Get_predData(data); torch::autograd::Variable fl = model.forward(predData); torch::autograd::Variable result = torch::squeeze(fl); torch::autograd::Variable rt = result.argmax(0); std::cout<<rt.size(0)<<std::endl; std::cout<<rt.size(1)<<std::endl; write2Txt(rt); } int main() { unet net(7,2); trainConvNet(net); return 0; }