如何使用神经网络获取列表第n个元素并实现对应自定义层

你提到的第一种动态调整权重的方案不推荐使用：训练过程中动态修改参数值会打断计算图的梯度传递，其次batch场景下不同样本的k值不同，无法给所有样本生成统一的权重矩阵，实际落地会遇到很多问题。

你第二种思路的方向是对的，不需要设置可训练参数，也不用做权重修改操作，直接用框架原生的one-hot算子和逐元素乘就能实现，全程支持自动求导，可以直接嵌入到现有深度网络中使用。

PyTorch 实现示例（无参数版本，最推荐） #

import torch
import torch.nn as nn

class KthElementSelector(nn.Module):
    def __init__(self, input_dim: int):
        super().__init__()
        self.input_dim = input_dim
        # 无任何可训练参数，不需要做权重冻结操作

    def forward(self, x: torch.Tensor, k: torch.Tensor) -> torch.Tensor:
        # 输入参数说明：
        # x: 前层网络输出，shape 为 [batch_size, input_dim]
        # k: 索引张量，shape 为 [batch_size]，每个元素为对应样本要选择的索引值
        # 输出为每个样本的第k个元素，shape 为 [batch_size]
        
        # 生成对应k的one-hot张量
        one_hot_k = nn.functional.one_hot(k, num_classes=self.input_dim).to(x.dtype)
        # 逐元素相乘后按最后一维求和，得到选中的第k个元素
        selected_element = torch.sum(x * one_hot_k, dim=-1)
        return selected_element

该实现的优势：

• 无任何可训练参数，不需要处理权重冻结、动态修改的逻辑，不会破坏计算图一致性
• 支持批量训练，batch内每个样本可以传入不同的k值，计算效率和原生算子一致
• 所有操作均为框架原生支持的可导算子，梯度可以正常回传到前层网络，支持端到端训练

如果需要适配更高维度的输入（比如形状为[batch_size, seq_len, hidden_dim]的序列特征，要选择第k个时间步的输出），只需要调整one-hot的广播维度和求和维度即可。

基于固定权重的实现（适配你第二种思路的要求） #

如果你一定要用冻结权重相乘的逻辑实现，可以参考以下代码，和上述无参数版本逻辑完全等价，仅多了无实际作用的冻结权重占位：

import torch
import torch.nn as nn

class KthElementSelectorFixedWeight(nn.Module):
    def __init__(self, input_dim: int):
        super().__init__()
        self.input_dim = input_dim
        # 初始化全0权重，设置为不可训练（冻结）
        self.fixed_weight = nn.Parameter(torch.zeros(input_dim), requires_grad=False)

    def forward(self, x: torch.Tensor, k: torch.Tensor) -> torch.Tensor:
        one_hot_k = nn.functional.one_hot(k, num_classes=self.input_dim).to(x.dtype)
        # 逻辑和无参数版本完全一致，仅复用了冻结权重的 dtype 设备属性，无实际计算作用
        selected_element = torch.sum(x * one_hot_k * self.fixed_weight.add(1), dim=-1)
        return selected_element

内容的提问来源于stack exchange，提问作者prakash gawas