CodeMirror/waifu2x
1
0
Fork 0
mirror of synced 2024-05-19 04:12:19 +12:00
Code Issues Releases Wiki Activity

Add -force_cudnn option to tools/benchmark.lua

Browse source
This commit is contained in:
nagadomi 2016-06-12 22:03:19 +09:00
parent 67d36a1220
commit 6cca2f8488
1 changed files with 13 additions and 14 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

27
tools/benchmark.lua
View file

@ -38,6 +38,7 @@ cmd:option("-tta", 0, 'use tta')
cmd:option("-tta_level", 8, 'tta level')
cmd:option("-crop_size", 128, 'patch size per process')
cmd:option("-batch_size", 1, 'batch_size')
cmd:option("-force_cudnn", 0, 'use cuDNN backend')
local function to_bool(settings, name)
if settings[name] == 1 then
@ -50,8 +51,9 @@ local opt = cmd:parse(arg)
torch.setdefaulttensortype('torch.FloatTensor')
if cudnn then
cudnn.fastest = true
cudnn.benchmark = false
cudnn.benchmark = true
end
to_bool(opt, "force_cudnn")
to_bool(opt, "save_all")
to_bool(opt, "tta")
if opt.save_all then
@ -341,17 +343,14 @@ local function load_data(test_dir)
end
return test_x
end
function load_model(filename)
return torch.load(filename, "ascii")
end
function load_noise_scale_model(model_dir, noise_level)
function load_noise_scale_model(model_dir, noise_level, force_cudnn)
local f = path.join(model_dir, string.format("noise%d_scale2.0x_model.t7", opt.noise_level))
local s1, noise_scale = pcall(load_model, f)
local s1, noise_scale = pcall(w2nn.load_model, f, force_cudnn)
local model = {}
if not s1 then
f = path.join(model_dir, string.format("noise%d_model.t7", opt.noise_level))
local noise
s1, noise = pcall(load_model, f)
s1, noise = pcall(w2nn.load_model, f, force_cudnn)
if not s1 then
model.noise_model = nil
print(model_dir .. "'s noise model is not found. benchmark will use only scale model.")
@ -360,7 +359,7 @@ function load_noise_scale_model(model_dir, noise_level)
end
f = path.join(model_dir, "scale2.0x_model.t7")
local scale
s1, scale = pcall(load_model, f)
s1, scale = pcall(w2nn.load_model, f, force_cudnn)
if not s1 then
return nil
end
@ -377,8 +376,8 @@ end
if opt.method == "scale" then
local f1 = path.join(opt.model1_dir, "scale2.0x_model.t7")
local f2 = path.join(opt.model2_dir, "scale2.0x_model.t7")
local s1, model1 = pcall(load_model, f1)
local s2, model2 = pcall(load_model, f2)
local s1, model1 = pcall(w2nn.load_model, f1, opt.force_cudnn)
local s2, model2 = pcall(w2nn.load_model, f2, opt.force_cudnn)
if not s1 then
error("Load error: " .. f1)
end
@ -390,8 +389,8 @@ if opt.method == "scale" then
elseif opt.method == "noise" then
local f1 = path.join(opt.model1_dir, string.format("noise%d_model.t7", opt.noise_level))
local f2 = path.join(opt.model2_dir, string.format("noise%d_model.t7", opt.noise_level))
local s1, model1 = pcall(load_model, f1)
local s2, model2 = pcall(load_model, f2)
local s1, model1 = pcall(w2nn.load_model, f1, opt.force_cudnn)
local s2, model2 = pcall(w2nn.load_model, f2, opt.force_cudnn)
if not s1 then
error("Load error: " .. f1)
end
@ -402,9 +401,9 @@ elseif opt.method == "noise" then
benchmark(opt, test_x, transform_jpeg, model1, model2)
elseif opt.method == "noise_scale" then
local model2 = nil
local model1 = load_noise_scale_model(opt.model1_dir, opt.noise_level)
local model1 = load_noise_scale_model(opt.model1_dir, opt.noise_level, opt.force_cudnn)
if opt.model2_dir:len() > 0 then
model2 = load_noise_scale_model(opt.model2_dir, opt.noise_level)
model2 = load_noise_scale_model(opt.model2_dir, opt.noise_level, opt.force_cudnn)
end
local test_x = load_data(opt.dir)
benchmark(opt, test_x, transform_scale_jpeg, model1, model2)