Python PyTorch LowRankMixtureCrossNet用法及代碼示例

本文簡要介紹python語言中 torchrec.modules.crossnet.LowRankMixtureCrossNet 的用法。

用法:

class torchrec.modules.crossnet.LowRankMixtureCrossNet(in_features: int, num_layers: int, num_experts: int = 1, low_rank: int = 1, activation: typing.Union[torch.nn.modules.module.Module, typing.Callable[[torch.Tensor], torch.Tensor]] = <built-in method relu of type object>)

參數：

• in_features(int) -輸入的維度。

• num_layers(int) -模塊中的層數。

• low_rank(int) -交叉矩陣的秩設置(默認 = 0)。值必須始終 >= 0

• activation(聯盟[火炬.nn.模塊,可調用[[torch.Tensor],torch.Tensor]]) - 非線性激活函數，用於定義專家。默認為 relu。

基礎：torch.nn.modules.module.Module

Low Rank Mixture Cross Net 是來自 paper 的 DCN V2 實現：

LowRankMixtureCrossNet 將每層的可學習交叉參數定義為低秩矩陣 \((N*r)\) 以及專家的混合。與LowRankCrossNet相比，該模塊利用了\(K\)專家，而不是依賴一位專家來學習特征組合；每個學習特征在不同子空間中相互作用，並使用取決於輸入 \(x\) 的門控機製自適應地組合學習到的交叉。

在每一層 l 上，張量轉換為：

\[x_{l+1} = MoE({expert_i : i \in K_{experts}}) + x_l\]

每個 \(expert_i\) 定義為：

\[expert_i = x_0 * (U_{li} \dot g(C_{li} \dot g(V_{li} \dot x_l)) + b_l)\]

其中 \(U_{li} (N, r)\) 、 \(C_{li} (r, r)\) 和 \(V_{li} (r, N)\) 是低秩矩陣，\(*\) 表示逐元素乘法，\(x\) 表示矩陣乘法，\(g()\) 是非線性激活函數。

num_expert 為 1 時，將跳過門評估和 MOE 以節省計算。

例子：

batch_size = 3
num_layers = 2
in_features = 10
input = torch.randn(batch_size, in_features)
dcn = LowRankCrossNet(num_layers=num_layers, num_experts=5, low_rank=3)
output = dcn(input)