Add -gpu option
This commit is contained in:
parent
16731fb634
commit
50fd999c38
|
@ -1,6 +1,7 @@
|
||||||
require 'xlua'
|
require 'xlua'
|
||||||
require 'pl'
|
require 'pl'
|
||||||
require 'trepl'
|
require 'trepl'
|
||||||
|
require 'cutorch'
|
||||||
|
|
||||||
-- global settings
|
-- global settings
|
||||||
|
|
||||||
|
@ -63,6 +64,7 @@ cmd:option("-oracle_drop_rate", 0.5, '')
|
||||||
cmd:option("-learning_rate_decay", 3.0e-7, 'learning rate decay (learning_rate * 1/(1+num_of_data*patches*epoch))')
|
cmd:option("-learning_rate_decay", 3.0e-7, 'learning rate decay (learning_rate * 1/(1+num_of_data*patches*epoch))')
|
||||||
cmd:option("-resume", "", 'resume model file')
|
cmd:option("-resume", "", 'resume model file')
|
||||||
cmd:option("-name", "user", 'model name for user method')
|
cmd:option("-name", "user", 'model name for user method')
|
||||||
|
cmd:option("-gpu", 1, 'Device ID')
|
||||||
|
|
||||||
local function to_bool(settings, name)
|
local function to_bool(settings, name)
|
||||||
if settings[name] == 1 then
|
if settings[name] == 1 then
|
||||||
|
@ -152,4 +154,6 @@ end
|
||||||
settings.images = string.format("%s/images.t7", settings.data_dir)
|
settings.images = string.format("%s/images.t7", settings.data_dir)
|
||||||
settings.image_list = string.format("%s/image_list.txt", settings.data_dir)
|
settings.image_list = string.format("%s/image_list.txt", settings.data_dir)
|
||||||
|
|
||||||
|
cutorch.setDevice(opt.gpu)
|
||||||
|
|
||||||
return settings
|
return settings
|
||||||
|
|
|
@ -267,6 +267,7 @@ local function waifu2x()
|
||||||
cmd:option("-tta_level", 8, 'TTA level (2|4|8). A higher value makes better quality output but slow')
|
cmd:option("-tta_level", 8, 'TTA level (2|4|8). A higher value makes better quality output but slow')
|
||||||
cmd:option("-force_cudnn", 0, 'use cuDNN backend (0|1)')
|
cmd:option("-force_cudnn", 0, 'use cuDNN backend (0|1)')
|
||||||
cmd:option("-q", 0, 'quiet (0|1)')
|
cmd:option("-q", 0, 'quiet (0|1)')
|
||||||
|
cmd:option("-gpu", 1, 'Device ID')
|
||||||
|
|
||||||
local opt = cmd:parse(arg)
|
local opt = cmd:parse(arg)
|
||||||
if opt.method:len() > 0 then
|
if opt.method:len() > 0 then
|
||||||
|
@ -292,5 +293,6 @@ local function waifu2x()
|
||||||
else
|
else
|
||||||
convert_frames(opt)
|
convert_frames(opt)
|
||||||
end
|
end
|
||||||
|
cutorch.setDevice(opt.gpu)
|
||||||
end
|
end
|
||||||
waifu2x()
|
waifu2x()
|
||||||
|
|
Loading…
Reference in a new issue