122 lines
3.3 KiB
Lua
122 lines
3.3 KiB
Lua
-- adapted from https://github.com/marcan/cl-waifu2x
|
|
require 'pl'
|
|
local __FILE__ = (function() return string.gsub(debug.getinfo(2, 'S').source, "^@", "") end)()
|
|
package.path = path.join(path.dirname(__FILE__), "..", "lib", "?.lua;") .. package.path
|
|
require 'w2nn'
|
|
local cjson = require "cjson"
|
|
|
|
local function meta_data(model, model_path)
|
|
local meta = {}
|
|
for k, v in pairs(model) do
|
|
if k:match("w2nn_") then
|
|
meta[k:gsub("w2nn_", "")] = v
|
|
end
|
|
end
|
|
|
|
modtime = file.modified_time(model_path)
|
|
utc_date = Date('utc')
|
|
utc_date:set(modtime)
|
|
meta["created_at"] = tostring(utc_date)
|
|
|
|
return meta
|
|
end
|
|
local function includes(s, a)
|
|
for i = 1, #a do
|
|
if s == a[i] then
|
|
return true
|
|
end
|
|
end
|
|
return false
|
|
end
|
|
local function get_bias(mod)
|
|
if mod.bias then
|
|
return mod.bias:float()
|
|
else
|
|
-- no bias
|
|
return torch.FloatTensor(mod.nOutputPlane):zero()
|
|
end
|
|
end
|
|
local function export_weight(jmodules, seq)
|
|
local convolutions = {"nn.SpatialConvolutionMM",
|
|
"cudnn.SpatialConvolution",
|
|
"cudnn.SpatialDilatedConvolution",
|
|
"nn.SpatialFullConvolution",
|
|
"nn.SpatialDilatedConvolution",
|
|
"cudnn.SpatialFullConvolution"
|
|
}
|
|
for k = 1, #seq.modules do
|
|
local mod = seq.modules[k]
|
|
local name = torch.typename(mod)
|
|
if name == "nn.Sequential" or name == "nn.ConcatTable" then
|
|
export_weight(jmodules, mod)
|
|
elseif name == "nn.Linear" then
|
|
local weight = torch.totable(mod.weight:float())
|
|
local jmod = {
|
|
class_name = name,
|
|
nInputPlane = mod.weight:size(2),
|
|
nOutputPlane = mod.weight:size(1),
|
|
bias = torch.totable(get_bias(mod)),
|
|
weight = weight
|
|
}
|
|
table.insert(jmodules, jmod)
|
|
elseif includes(name, convolutions) then
|
|
local weight = mod.weight:float()
|
|
if name:match("FullConvolution") then
|
|
weight = torch.totable(weight:reshape(mod.nInputPlane, mod.nOutputPlane, mod.kH, mod.kW))
|
|
else
|
|
weight = torch.totable(weight:reshape(mod.nOutputPlane, mod.nInputPlane, mod.kH, mod.kW))
|
|
end
|
|
local jmod = {
|
|
class_name = name,
|
|
kW = mod.kW,
|
|
kH = mod.kH,
|
|
dH = mod.dH,
|
|
dW = mod.dW,
|
|
padW = mod.padW,
|
|
padH = mod.padH,
|
|
dilationW = mod.dilationW,
|
|
dilationH = mod.dilationH,
|
|
nInputPlane = mod.nInputPlane,
|
|
nOutputPlane = mod.nOutputPlane,
|
|
bias = torch.totable(get_bias(mod)),
|
|
weight = weight
|
|
}
|
|
table.insert(jmodules, jmod)
|
|
end
|
|
end
|
|
end
|
|
local function export(model, model_path, output)
|
|
local jmodules = {}
|
|
local model_config = meta_data(model, model_path)
|
|
local first_layer = true
|
|
|
|
print(model_config)
|
|
print(model)
|
|
|
|
export_weight(jmodules, model)
|
|
jmodules[1]["model_config"] = model_config
|
|
|
|
local fp = io.open(output, "w")
|
|
if not fp then
|
|
error("IO Error: " .. output)
|
|
end
|
|
fp:write(cjson.encode(jmodules))
|
|
fp:close()
|
|
end
|
|
|
|
local cmd = torch.CmdLine()
|
|
cmd:text()
|
|
cmd:text("waifu2x export model")
|
|
cmd:text("Options:")
|
|
cmd:option("-i", "input.t7", 'Specify the input torch model')
|
|
cmd:option("-o", "output.json", 'Specify the output json file')
|
|
cmd:option("-iformat", "ascii", 'Specify the input format (ascii|binary)')
|
|
|
|
local opt = cmd:parse(arg)
|
|
if not path.isfile(opt.i) then
|
|
cmd:help()
|
|
os.exit(-1)
|
|
end
|
|
local model = torch.load(opt.i, opt.iformat)
|
|
export(model, opt.i, opt.o)
|