refactor tools/export_model.lua

tools/export_model.lua

@@ -4,21 +4,45 @@ package.path = path.join(path.dirname(__FILE__), "..", "lib", "?.lua;") .. packa
 require 'w2nn'
 local cjson = require "cjson"
 
-local model = torch.load(arg[1], "ascii")
-
-local jmodules = {}
-local modules = model:findModules("nn.SpatialConvolutionMM")
-for i = 1, #modules, 1 do
-   local module = modules[i]
-   local jmod = {
-      kW = module.kW,
-      kH = module.kH,
-      nInputPlane = module.nInputPlane,
-      nOutputPlane = module.nOutputPlane,
-      bias = torch.totable(module.bias:float()),
-      weight = torch.totable(module.weight:float():reshape(module.nOutputPlane, module.nInputPlane, module.kW, module.kH))
-   }
-   table.insert(jmodules, jmod)
+function export(model, output)
+   local jmodules = {}
+   local modules = model:findModules("nn.SpatialConvolutionMM")
+   if #modules == 0 then
+      -- cudnn model
+      modules = model:findModules("cudnn.SpatialConvolution")
+   end
+   for i = 1, #modules, 1 do
+      local module = modules[i]
+      local jmod = {
+	 kW = module.kW,
+	 kH = module.kH,
+	 nInputPlane = module.nInputPlane,
+	 nOutputPlane = module.nOutputPlane,
+	 bias = torch.totable(module.bias:float()),
+	 weight = torch.totable(module.weight:float():reshape(module.nOutputPlane, module.nInputPlane, module.kW, module.kH))
+      }
+      table.insert(jmodules, jmod)
+   end
+   local fp = io.open(output, "w")
+   if not fp then
+      error("IO Error: " .. output)
+   end
+   fp:write(cjson.encode(jmodules))
+   fp:close()
 end
 
-io.write(cjson.encode(jmodules))
+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.o)