从 DBPN(Pytorch) 导出模型训练参数

在超分辨重建《DBPN-Pytorch-master》包中有已经训练好的模型,这个DBPN好象是“2018超分辨率 8 倍的冠军”。

所以风格迁移就先放一放,把这个DBPN超分辨用C++实行一下。

先安装好 Pytorch

运行 python eval.py 或者 eval_gan.py,就可以用该超分辨了。

从 DBPN(Pytorch) 导出模型训练参数_第1张图片输入图

从 DBPN(Pytorch) 导出模型训练参数_第2张图片eval_gan.py 4倍  运行时间: 348.8485 秒(CPU)

8倍非常费内存,输入图缩小一些

输入图

从 DBPN(Pytorch) 导出模型训练参数_第3张图片eval.py 8倍  运行时间: 157.8425 秒(CPU)

 

导出模型中所有的训练参数(在载入模型参数后处加入代码,按caffe格式):

#---------------------------------------->
# 移除字符串 `module.` 
model.load_state_dict({k.replace('module.',''):v for k,v in torch.load(opt.model, map_location=lambda storage, loc: storage).items()})
#----------------------------------------<
print('预先训练的SR模型已加载.')

#------------------------------------------------------------->
#导出模型中所有的可训练参数
#遍历named_parameters()中的所有的参数,只打印那些param.requires_grad=True的变量。
fp = open('PIRM2018_region2.pth_x4model.txt', 'w')
for name, param in model.named_parameters():
    print(name)
    fp.write(name) #层名
    fp.write('\n')


    if param.requires_grad:

        fp.write('{\n')

        #维度顺序应该和torch一样吧,不用转换
        v_1d=param.view(-1) #直接转为一维
        fp.write('  blobs {\n')

        for vv in v_1d:

            fp.write('%8f' % vv) # '    data: %8f'

            fp.write('\n')

        fp.write('    shape {\n')

        for s in param.size():

            fp.write('      dim: ' + str(s))# 维度

            fp.write('\n')

        fp.write('    }\n')
        fp.write('  }\n')
        fp.write('}\n')

fp.close()
os.system("pause")

#-------------------------------------------------------------<

已经成为一个154M的文件了:

feat0.conv.weight
{
  blobs {
0.281878
-0.380647
0.184301
省略...
0.051228
-0.140397
-0.387719
    shape {
      dim: 256
      dim: 3
      dim: 3
      dim: 3
    }
  }
}
feat0.conv.bias
{
  blobs {
-0.019276
-0.055616
0.008202
省略...
-0.011527
-0.050689
0.050540
    shape {
      dim: 256
    }
  }
}
feat0.act.weight
{
  blobs {
0.780346
    shape {
      dim: 1
    }
  }
}
feat1.conv.weight
{
  blobs {
0.060832
-0.090450
0.080900
省略...

先这样吧

 

你可能感兴趣的:(超分辨率重建,导出pytorch模型参数)