add cudnn converter
This commit is contained in:
parent
da786e15ba
commit
21ea5dd858
|
@ -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,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
tools/cudnn2cunn.lua
Normal file
47
tools/cudnn2cunn.lua
Normal file
|
@ -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
tools/cunn2cudnn.lua
Normal file
47
tools/cunn2cudnn.lua
Normal file
|
@ -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)
|
Loading…
Reference in a new issue