From 609d4f862202c43241fdab06b0795db6281ca3f7 Mon Sep 17 00:00:00 2001 From: Vsevolod Stakhov Date: Sun, 17 Sep 2017 10:04:59 +0100 Subject: [PATCH] [Feature] Allow to specify number of threads for ANN learning --- lualib/lua_nn.lua | 1 + src/plugins/lua/fann_redis.lua | 4 ++++ 2 files changed, 5 insertions(+) diff --git a/lualib/lua_nn.lua b/lualib/lua_nn.lua index 0e8977d374..d0d2d5265d 100644 --- a/lualib/lua_nn.lua +++ b/lualib/lua_nn.lua @@ -22,6 +22,7 @@ local lua_nn_models = {} if rspamd_config:has_torch() then torch = require "torch" + torch.setnumthreads(1) end if torch then diff --git a/src/plugins/lua/fann_redis.lua b/src/plugins/lua/fann_redis.lua index 2751b5d791..ab3da00350 100644 --- a/src/plugins/lua/fann_redis.lua +++ b/src/plugins/lua/fann_redis.lua @@ -47,6 +47,7 @@ local default_options = { mse = 0.001, autotrain = true, train_prob = 1.0, + learn_threads = 1, }, use_settings = false, per_user = false, @@ -781,6 +782,9 @@ local function train_fann(rule, _, ev_base, elt, worker) dataset.size = function() return #dataset end local function train_torch() + if rule.train.learn_threads > 1 then + torch.setnumthreads(rule.train.learn_threads) + end local criterion = nn.MSECriterion() local trainer = nn.StochasticGradient(fanns[elt].fann_train, criterion) -- 2.47.3