Browse Source

Add support for MSE loss

nagadomi 8 years ago
parent
commit
0eccbc6555
2 changed files with 3 additions and 1 deletions
  1. 1 1
      lib/settings.lua
  2. 2 0
      train.lua

+ 1 - 1
lib/settings.lua

@@ -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 - 0
train.lua

@@ -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