[linear-ridge-svd]使用 SPU 实现线性回归模型

[Candicepan]

此 ISSUE 为 隐语开源共建计划（SecretFlow Open Source Contribution Plan，简称 SF OSCP）第二期 #249 子任务，若对该任务感兴趣，欢迎在 #249 提交你的不同设计思路。
也欢迎其他开发者点击这里，选择其他你感兴趣的任务，热烈欢迎社区开发者参与共建～

设计思路

by：magic-hya

1. 参考sklearn的ridge算法，在jax上实现
2. 实现后与明文sklearn ridge算法结果做对比

任务介绍

• 任务名称：使用 SPU 实现线性回归模型
• 技术方向：SPU/SML
• 任务难度：热身🌟

详细要求

• 安全性（尽量少 reveal）
• 功能性：实现算法的基本功能，包括：
  • 支持 fit（至少一种 eg：sgd, lbfgs, svd...），predict
  • 提供配置是否需要拟合 bias 项，是否需要正则等基本参数
  • 支持起码一种正则方式（L2, L1, elastic...）
• 收敛性：包含 simulator 跑出的实验数据并且证明收敛性/准确性（最好有与明文 sklearn 结果的对比，可详见参考 PR）
• 代码规范：Python 代码需要使用 black+isort 进行格式化（流水线包含代码规范检查卡点）
• 提交说明：关联该 isuue 并提交代码至 https://github.com/secretflow/spu/tree/main/sml
• 特殊说明：若某个特性有特殊的限制，如需要 FM128，需要更多 fxp 等需要在注释文档中明确说明

能力要求

• 熟悉经典的机器学习算法
• 熟悉 JAX 或 NumPy，可以使用 NumPy 实现算法

操作说明

• 操作指引：https://www.secretflow.org.cn/docs/spu/latest/en-US/getting_started/tutorials/develop_your_first_mpc_application
• 可参考明文实现：
  • linear_model/_base
  • linear_model/_ridge
• 范例
• 参考 PR

认领说明

本任务可有多种实现方式，故支持一个任务有多位开发者进行认领，请在认领任务后，在该 issue 下 comment 你的具体设计思路。

设计思路说明：简单说明计划使用什么算法 or 什么优化器实现任务需求即可

当同一任务有多位开发者提交设计思路时：

• 若多位开发者的设计思路均类似，则将按照 comment 的时间 assign 给第一位 comment 的开发者；
• 若多位开发者的设计思路均不同，则该 issue 将会拆分为以不同设计思路进行实现的子 issue，并 assign 给对应开发者。

开发须知

以下部分代码请必须增加代码注释，对对应代码模块进行说明，包括：

• __init__函数的超参数含义
• fit 的具体算法实现说明

[magic-hya]

magic-hya Give it to me

[deadlywing]

hello，，请问您当前实现了哪种solver去实现ridge regression？

[magic-hya]

hello，，请问您当前实现了哪种solver去实现ridge regression？

hello, 我是参考sklearn里面的做法实现，里面实现的方法很多，我也打算逐步实现里面的内容，丰富算法，目前第一个版本先实现简单的svd(奇异值分解)

[deadlywing]

实现一个solver以后就可以直接先发pr哈，不然这个pr的持续时间就太长了～

如果您还有实现其他solver的打算，可以再发起其他pr～

Thanks

[magic-hya]

大概明后天吧，我会把我第一版的代码提交pr

[deadlywing]

hello，，请问您当前实现了哪种solver去实现ridge regression？

hello, 我是参考sklearn里面的做法实现，里面实现的方法很多，我也打算逐步实现里面的内容，丰富算法，目前第一个版本先实现简单的svg

sorry，，svg？是svd还是sag呀？

[magic-hya]

hello，，请问您当前实现了哪种solver去实现ridge regression？

hello, 我是参考sklearn里面的做法实现，里面实现的方法很多，我也打算逐步实现里面的内容，丰富算法，目前第一个版本先实现简单的svg

sorry，，svg？是svd还是sag呀？

sorry，svd 奇异值分解

[magic-hya]

目前本地测试没有问题，但是执行emul和test会报错。
执行python sml/linear_model/ridge_emul.py
错误：
Traceback (most recent call last): File "sml/linear_model/ridge_emul.py", line 31, in <module> import sml.utils.emulation as emulation File "/data/linear-ridge/spuxx/sml/utils/emulation.py", line 85, in <module> class Emulator: File "/data/linear-ridge/spuxx/sml/utils/emulation.py", line 163, in Emulator copts=spu_pb2.CompilerOptions(), AttributeError: module 'spu.spu_pb2' has no attribute 'CompilerOptions'

执行python sml/linear_model/ridge_emul.py
错误：
Traceback (most recent call last): File "sml/linear_model/ridge_test.py", line 49, in test_simple result = spsim.sim_jax(sim, proc)(x1, x2, y) File "/usr/local/lib/python3.8/site-packages/spu/utils/simulation.py", line 152, in wrapper executable, output = spu_fe.compile( File "/usr/local/lib/python3.8/site-packages/spu/utils/frontend.py", line 112, in compile ir_text, output = _jax_compilation(fn, static_argnums, args, kwargs) File "/usr/local/lib/python3.8/site-packages/cachetools/__init__.py", line 737, in wrapper v = func(*args, **kwargs) File "/usr/local/lib/python3.8/site-packages/spu/utils/frontend.py", line 41, in _jax_compilation cfn, output = jax.xla_computation( File "sml/linear_model/ridge_test.py", line 42, in proc model.fit(x, y).predict_proba(x) File "/data/linear-ridge/spuxx/sml/linear_model/ridge.py", line 78, in fit self.coef = _solve_svd(x, y, alpha) File "/data/linear-ridge/spuxx/sml/linear_model/ridge.py", line 105, in _solve_svd s_nnz = s[idx][:, jnp.newaxis] File "/usr/local/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py", line 3907, in _rewriting_take treedef, static_idx, dynamic_idx = _split_index_for_jit(idx, arr.shape) File "/usr/local/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py", line 3985, in _split_index_for_jit idx = _expand_bool_indices(idx, shape) File "/usr/local/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py", line 4299, in _expand_bool_indices raise errors.NonConcreteBooleanIndexError(abstract_i) jax._src.errors.NonConcreteBooleanIndexError: Array boolean indices must be concrete; got ShapedArray(bool[10])

详细代码见附件ridge.zip
ridge.zip

[tpppppub]

请用 bazel test/run 的方式运行 test 和 emul

[magic-hya]

目前遇到问题，无法通过测试，执行测试报错，主要代码为以下函数：
`

def _solve_svd(x, y, alpha):

 U, s, Vt = jsci.linalg.svd(x, full_matrices=False)
 s_nnz = s[:, jnp.newaxis]
 UTy = jnp.dot(U.T, y)
 d = s_nnz / (s_nnz**2 + alpha)
 d_UT_y = d * UTy
 return jnp.dot(Vt.T, d_UT_y).T`

错误如下：jax._src.traceback_util.UnfilteredStackTrace: NotImplementedError: MLIR translation rule for primitive 'eigh' not found for platform interpreter

[deadlywing]

spu暂时不支持svd，eigh等于特征值相关的api，近期可能会commit一个utils文件，里面会有一个简单的svd实现。您可以：

1. 使用其他的解法
2. 自己用jax实现一个svd

[tpppppub]

@magic-hya 可以参考下这个 issue #213

[magic-hya]

spu暂时不支持svd，eigh等于特征值相关的api，近期可能会commit一个utils文件，里面会有一个简单的svd实现。您可以：

1. 使用其他的解法
2. 自己用jax实现一个svd

使用cholesky实现了，也跑通了emul和test测试，跟sklearn结果误差在0.01-0.02之间，我先提交个pr，然后看下还有什么要改的