From: Vsevolod Stakhov <vsevolod@highsecure.ru>
Date: Sun, 17 Sep 2017 09:04:59 +0000 (+0100)
Subject: [Feature] Allow to specify number of threads for ANN learning
X-Git-Tag: 1.7.0~629
X-Git-Url: http://git.ipfire.org/gitweb.cgi?a=commitdiff_plain;h=609d4f862202c43241fdab06b0795db6281ca3f7;p=thirdparty%2Frspamd.git

[Feature] Allow to specify number of threads for ANN learning
---

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)