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

Add support for MSE loss

Browse source
This commit is contained in:
nagadomi 2017-01-09 13:00:00 +09:00
parent 6051dad793
commit 0eccbc6555
2 changed files with 3 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
lib/settings.lua
View file

@ -75,7 +75,7 @@ cmd:option("-learning_rate_decay", 3.0e-7, 'learning rate decay (learning_rate *
cmd:option("-resume", "", 'resume model file')
cmd:option("-name", "user", 'model name for user method')
cmd:option("-gpu", 1, 'Device ID')
cmd:option("-loss", "huber", 'loss function (huber|l1)')
cmd:option("-loss", "huber", 'loss function (huber|l1|mse)')
local function to_bool(settings, name)
if settings[name] == 1 then

2
train.lua
View file

@ -317,6 +317,8 @@ local function create_criterion(model)
end
elseif settings.loss == "l1" then
return w2nn.L1Criterion():cuda()
elseif settings.loss == "mse" then
return w2nn.ClippedMSECriterion(0, 1.0):cuda()
else
error("unsupported loss .." .. settings.loss)
end