aboutsummaryrefslogtreecommitdiff
path: root/nerv/layer/mse.lua
diff options
context:
space:
mode:
Diffstat (limited to 'nerv/layer/mse.lua')
-rw-r--r--nerv/layer/mse.lua17
1 files changed, 11 insertions, 6 deletions
diff --git a/nerv/layer/mse.lua b/nerv/layer/mse.lua
index 0ee3080..1c218d0 100644
--- a/nerv/layer/mse.lua
+++ b/nerv/layer/mse.lua
@@ -5,6 +5,11 @@ function MSELayer:__init(id, global_conf, layer_conf)
self.dim_in = layer_conf.dim_in
self.dim_out = layer_conf.dim_out
self.gconf = global_conf
+ if self.gconf.use_cpu then
+ self.mat_type = self.gconf.mmat_type
+ else
+ self.mat_type = self.gconf.cumat_type
+ end
self:check_dim_len(2, -1)
end
@@ -15,15 +20,15 @@ function MSELayer:init(batch_size)
self.scale = 1 / self.dim_in[1]
self.total_mse = 0.0
self.total_frames = 0
- self.mse = self.gconf.cumat_type(batch_size, self.dim_in[1])
- self.mse_sum = self.gconf.cumat_type(batch_size, 1)
+ self.mse = self.mat_type(batch_size, self.dim_in[1])
+ self.mse_sum = self.mat_type(batch_size, 1)
self.diff = self.mse:create()
end
function MSELayer:batch_resize(batch_size)
if self.mse:nrow() ~= batch_resize then
- self.mse = self.gconf.cumat_type(batch_size, self.dim_in[1])
- self.mse_sum = self.gconf.cumat_type(batch_size, 1)
+ self.mse = self.mat_type(batch_size, self.dim_in[1])
+ self.mse_sum = self.mat_type(batch_size, 1)
self.diff = self.mse:create()
end
end
@@ -36,11 +41,11 @@ function MSELayer:propagate(input, output)
local mse = self.mse
local mse_sum = self.mse_sum
mse:add(input[1], input[2], 1.0, -1.0)
- self.diff:copy_fromd(mse)
+ self.diff:copy_from(mse)
mse:mul_elem(mse, mse)
mse_sum:add(mse_sum, mse:rowsum(mse), 0.0, self.scale)
if output[1] ~= nil then
- output[1]:copy_fromd(mse_sum)
+ output[1]:copy_from(mse_sum)
end
self.total_mse = self.total_mse + mse_sum:colsum()[0][0]
self.total_frames = self.total_frames + mse_sum:nrow()