Forráskód Böngészése

add cudnn converter

nagadomi 9 éve
szülő
commit
21ea5dd858
4 módosított fájl, 97 hozzáadás és 34 törlés
  1. 0 34
      cudnn2cunn.lua
  2. 3 0
      lib/srcnn.lua
  3. 47 0
      tools/cudnn2cunn.lua
  4. 47 0
      tools/cunn2cudnn.lua

+ 0 - 34
cudnn2cunn.lua

@@ -1,34 +0,0 @@
-require 'cunn'
-require 'cudnn'
-require 'cutorch'
-require './lib/LeakyReLU'
-local srcnn = require 'lib/srcnn'
-
-local function cudnn2cunn(cudnn_model)
-   local cunn_model = srcnn.waifu2x("y")
-   local from_seq = cudnn_model:findModules("cudnn.SpatialConvolution")
-   local to_seq = cunn_model:findModules("nn.SpatialConvolutionMM")
-
-   for i = 1, #from_seq do
-      local from = from_seq[i]
-      local to = to_seq[i]
-      to.weight:copy(from.weight)
-      to.bias:copy(from.bias)
-   end
-   cunn_model:cuda()
-   cunn_model:evaluate()
-   return cunn_model
-end
-
-local cmd = torch.CmdLine()
-cmd:text()
-cmd:text("convert cudnn model to cunn model ")
-cmd:text("Options:")
-cmd:option("-model", "./model.t7", 'path of cudnn model file')
-cmd:option("-iformat", "ascii", 'input format')
-cmd:option("-oformat", "ascii", 'output format')
-
-local opt = cmd:parse(arg)
-local cudnn_model = torch.load(opt.model, opt.iformat)
-local cunn_model = cudnn2cunn(cudnn_model)
-torch.save(opt.model, cunn_model, opt.oformat)

+ 3 - 0
lib/srcnn.lua

@@ -3,6 +3,9 @@ require 'w2nn'
 -- ref: http://arxiv.org/abs/1502.01852
 -- ref: http://arxiv.org/abs/1501.00092
 local srcnn = {}
+function srcnn.channels(model)
+   return model:get(model:size() - 1).weight:size(1)
+end
 function srcnn.waifu2x_cunn(ch)
    local model = nn.Sequential()
    model:add(nn.SpatialConvolutionMM(ch, 32, 3, 3, 1, 1, 0, 0))

+ 47 - 0
tools/cudnn2cunn.lua

@@ -0,0 +1,47 @@
+local __FILE__ = (function() return string.gsub(debug.getinfo(2, 'S').source, "^@", "") end)()
+package.path = path.join(path.dirname(__FILE__), "..", "lib", "?.lua;") .. package.path
+require 'os'
+require 'pl'
+require 'torch'
+require 'cutorch'
+require 'cunn'
+require 'cudnn'
+require 'w2nn'
+local srcnn = require 'srcnn'
+
+local function cudnn2cunn(cunn_model)
+   local cudnn_model = srcnn.waifu2x_cunn(srcnn.channels(cunn_model))
+   local weight_from = cudnn_model:findModules("cudnn.SpatialConvolution")
+   local weight_to = cunn_model:findModules("nn.SpatialConvolutionMM")
+   
+   assert(#weight_from == #weight_to)
+   
+   for i = 1, #weight_from do
+      local from = weight_from[i]
+      local to = weight_to[i]
+      
+      to.weight:copy(from.weight)
+      to.bias:copy(from.bias)
+   end
+   cunn_model:cuda()
+   cunn_model:evaluate()
+   return cunn_model
+end
+
+local cmd = torch.CmdLine()
+cmd:text()
+cmd:text("waifu2x cudnn model to cunn model converter")
+cmd:text("Options:")
+cmd:option("-i", "", 'Specify the input cunn model')
+cmd:option("-o", "", 'Specify the output cudnn model')
+cmd:option("-iformat", "ascii", 'Specify the input format (ascii|binary)')
+cmd:option("-oformat", "ascii", 'Specify the output format (ascii|binary)')
+
+local opt = cmd:parse(arg)
+if not path.isfile(opt.i) then
+   cmd:help()
+   os.exit(-1)
+end
+local cudnn_model = torch.load(opt.i, opt.iformat)
+local cunn_model = cudnn2cunn(cudnn_model)
+torch.save(opt.o, cunn_model, opt.oformat)

+ 47 - 0
tools/cunn2cudnn.lua

@@ -0,0 +1,47 @@
+local __FILE__ = (function() return string.gsub(debug.getinfo(2, 'S').source, "^@", "") end)()
+package.path = path.join(path.dirname(__FILE__), "..", "lib", "?.lua;") .. package.path
+require 'os'
+require 'pl'
+require 'torch'
+require 'cutorch'
+require 'cunn'
+require 'cudnn'
+require 'w2nn'
+local srcnn = require 'srcnn'
+
+local function cunn2cudnn(cunn_model)
+   local cudnn_model = srcnn.waifu2x_cudnn(srcnn.channels(cunn_model))
+   local weight_from = cunn_model:findModules("nn.SpatialConvolutionMM")
+   local weight_to = cudnn_model:findModules("cudnn.SpatialConvolution")
+
+   assert(#weight_from == #weight_to)
+   
+   for i = 1, #weight_from do
+      local from = weight_from[i]
+      local to = weight_to[i]
+      
+      to.weight:copy(from.weight)
+      to.bias:copy(from.bias)
+   end
+   cudnn_model:cuda()
+   cudnn_model:evaluate()
+   return cudnn_model
+end
+
+local cmd = torch.CmdLine()
+cmd:text()
+cmd:text("waifu2x cunn model to cudnn model converter")
+cmd:text("Options:")
+cmd:option("-i", "", 'Specify the input cudnn model')
+cmd:option("-o", "", 'Specify the output cunn model')
+cmd:option("-iformat", "ascii", 'Specify the input format (ascii|binary)')
+cmd:option("-oformat", "ascii", 'Specify the output format (ascii|binary)')
+
+local opt = cmd:parse(arg)
+if not path.isfile(opt.i) then
+   cmd:help()
+   os.exit(-1)
+end
+local cunn_model = torch.load(opt.i, opt.iformat)
+local cudnn_model = cunn2cudnn(cunn_model)
+torch.save(opt.o, cudnn_model, opt.oformat)