diff options
Diffstat (limited to 'nerv/layer/mse.lua')
-rw-r--r-- | nerv/layer/mse.lua | 17 |
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() |