【Scala PyTorch深度学习】PyTorch On Scala系列课程 第十章 21 ：PyTorch微分【AI Infra 3.0】[PyTorch Scala 高校计算机硕士研一课程]

PyTorch Scala 高校计算机硕士研一课程神经常微分方程传统深度神经网络例如残差网络 (ResNets)通过离散的层序列处理输入。我们可以将ResNet块视为连续变换的欧拉离散化ht1htf(ht,θt)h**t1h**tf(h**t,θ**t)。这种观点自然引出一个问题我们能否连续地建模这种变换神经常微分方程 (Neural ODEs) 给出了肯定的答复它将网络深度定义为连续时间区间而非层数。连续深度框架不同于离散变换神经ODE使用常微分方程 (ODE) 建模隐藏状态 h(t)h(t) 随连续时间变量 tt的演变。其主要思想是使用神经网络 ff以权重 θθ为参数来定义隐藏状态随时间的变化率dh(t)dtf(h(t),t,θ)dtd**h(t)f(h(t),t,θ)在这里h(t)h(t) 表示时间 tt时的隐藏状态而 ff通常是一个标准神经网络例如一个MLP它以当前状态 h(t)h(t)、当前时间 tt和参数 θθ作为输入输出状态的变化率。将输入 z0z0即 h(t0)h(t0)转换为输出 z1z1即 h(t1)h(t1)的整体过程是通过在指定时间区间 [t0,t1][t0,t1] 上求解此ODE初始值问题得到的h(t1)h(t0)∫t0t1f(h(t),t,θ)dth(t1)h(t0)∫t0t1f(h(t),t,θ)d**t这个积分通过ODE求解器进行数值计算。神经网络 ff定义了向量场求解器模拟了隐藏状态通过该向量场从起始时间 t0t0 到结束时间 t1t1 的路径。神经ODE的优势这种连续的表述方式具有多项有益的特点训练时的内存效率标准反向传播需要存储每一层的激活值来计算梯度。对于层数多的网络或等同于ODE求解器正向传播中的许多步骤这会消耗大量内存。神经ODE采用伴随敏感度方法来计算梯度。该方法涉及逆向求解第二个相关的ODE。重要的是它计算参数 θθ和初始状态 h(t0)h(t0) 所需的梯度时内存使用量相对于“深度”或积分时间近似为常数。这使得训练具有复杂变换潜力的模型成为可能而无需承担存储中间状态带来的内存负担。自适应计算现代ODE求解器在积分过程中会自动调整步长。当动态 ff变化迅速时它们会采取较小的步长当动态平滑时则采取较大的步长。这意味着计算工作量可以适应所学函数的复杂性与ResNets等固定步长架构相比可能会带来更高的计算效率。处理不规则时间序列神经ODE天然适合建模连续过程和在不规则时间点采样的数据。模型可以通过将ODE积分到任意时间 tt来评估该点的隐藏状态。使用PyTorch实现实现神经ODE通常需要一个提供可微分ODE求解器的外部库。一个常用的选择是torchdiffeq。通常的工作流程包括定义动态函数创建一个标准torch.nn.Module来表示函数 f(h(t),t,θ)f(h(t),t,θ)。这个模块以当前状态h和时间t作为输入并返回计算出的导数dh/dt。importtorchimporttorch.nnas nnclassODEFuncextendsnn.Module:def__init__(hidden_dim:Int):super(ODEFunc,self).__init__()valnetnn.Sequential(nn.Linear(hidden_dim,hidden_dim),nn.Tanh(),nn.Linear(hidden_dim,hidden_dim),)defforward(t:Float,h:Tensor):// t当前时间标量// h当前隐藏状态张量// 返回 dh/dtreturnnet(h)使用ODE求解器使用torchdiffeq中的odeint等函数。该函数接收动态函数func、初始状态h0、要评估解的时间点t例如torch.tensor([t0, t1])以及可选的求解器参数。它返回在指定时间点计算出的隐藏状态。// 假设 torchdiffeq 已安装pip install torchdiffeqfrom torchdiffeqimportodeint_adjointas odeint// 使用伴随方法以节省内存// 示例用法valfuncODEFunc(hidden_dim20)valh0torch.randn(batch_size,20)// 初始状态valt_spantorch.tensor([0.0,1.0])// 从 t0 积分到 t1// 计算最终状态 h(t1)// odeint 通过伴随方法处理数值积分和梯度计算valh1odeint(func,h0,t_span)[-1]// 获取最后一个时间点 (t1) 的状态// h1 现在可以在后续层或损失函数中使用// 可以通过 h1.backward() 计算 func.parameters() 和 h0 的梯度注意odeint_adjoint的使用。该版本实现了内存高效的伴随反向传播方法。标准的odeint也可用但可能占用更多内存。伴随敏感度方法直接通过ODE求解器的操作进行反向传播可能会耗费大量计算资源和内存因为它需要存储求解器计算的所有中间状态。伴随方法提供了一种替代方案。它定义了伴随状态a(t)∂L∂h(t)a(t)∂h(t)∂L这表示最终损失 LL对隐藏状态 h(t)h(t) 的梯度。该伴随状态的演变由另一个逆时间从 t1t1 到 t0t0运行的ODE控制da(t)dt−a(t)T∂f(h(t),t,θ)∂hdtd**a(t)−a(t)T∂h∂f(h(t),t,θ)损失对参数 θθ的梯度可以通过逆时间积分另一个相关量来计算∂L∂θ∫t1t0a(t)T∂f(h(t),t,θ)∂θdt∂θ∂L∫t1t0a(t)T∂θ∂f(h(t),t,θ)d**t求解这些逆向ODE需要在反向传播过程中获取 h(t)h(t) 的值。然而可以通过再次求解原始正向ODE dh(t)dtf(h(t),t,θ)dtd**h(t)f(h(t),t,θ) 来实时重新计算这些值这次是从 h(t1)h(t1) 到 h(t0)h(t0) 逆向进行。这种重新计算避免了存储整个正向轨迹从而大大节省了内存通常将内存成本从 O(Nt)O(N**t) 降低到 O(1)O(1)其中 NtN**t是求解器步数。求解器选择与实际考量torchdiffeq等库提供了多种ODE求解器显式固定步长求解器欧拉法、中点法、RK4四阶龙格-库塔法。简单但为了稳定性和准确性可能需要较小的步长。自适应步长求解器Dormand-Prince (dopri5)、Adams方法。自动调整步长对于平滑问题通常更高效、更准确。dopri5常作为不错的默认选项。隐式求解器对于“刚性”ODE很有用因为如果步长不大显式方法会变得不稳定。它们通常在每一步都涉及方程求解计算开销可能更大。求解器的选择会影响准确性、稳定性和计算速度。它通常被视为一个需要调整的超参数。挑战计算成本尽管内存高效但求解ODE可能比固定数量的离散层计算速度慢特别是使用伴随方法进行反向传播时。数值稳定性求解器的选择以及所学动态函数 ff的性质会影响积分过程中的数值稳定性。如果 ff表现不佳仍然可能出现激活值爆炸或消失的问题。训练动态优化神经ODE有时比优化标准网络更具挑战性。神经ODE展现了深度学习与微分方程之间引人入胜的联系。它们提供了一种内存高效的方式来建模复杂、连续的变换并为涉及连续动态或不规则时间序列数据的问题提供了一种独特的工具从而扩展了PyTorch中可用的高级网络架构种类。元学习算法“标准深度学习模型在大型标记数据集上训练时通常表现优异。然而许多情况下需要用少量示例快速适应新任务这种场景称为少样本学习。元学习即“学习如何学习”提供了一个训练模型的体系使其能够有效泛化到数据有限的新任务。元学习算法不是学习如何很好地执行一个特定任务而是学习一个过程或一个初始化方法从而能够快速适应新的、相关任务。”主要说明如何在PyTorch中实现元学习算法特别是介绍一种流行且多功能的方案模型无关元学习MAML。元学习问题设置在典型的监督学习设置中我们有一个数据集 D{(xi,yi)}D{(x**i,y**i)}目标是学习一个由 θθ参数化的函数 fθf**θ使其在数据集上最小化损失 L(fθ(xi),yi)L(f**θ(x**i),y**i)。元学习重新定义了这个问题。我们假设存在任务分布 p(T)p(T)。在元训练期间我们从 p(T)p(T) 中采样批次任务 TiTi。对于每个任务 TiTi我们通常有一个小的支持集DisuppDisupp用于任务内部学习以及一个查询集DiqueryDiquery用于评估该任务的学习效果。目标是学习模型参数 θθ通常称为元参数使得模型能够利用新的、以前未见的任务 TnewTnew的支持集快速适应从而在其查询集 DnewqueryDnewquery上获得良好性能。模型无关元学习MAMLMAML 由 Finn 等人于 2017 年提出其目标是找到对任务变化敏感的元参数 θθ仅用少量梯度步长就能在小支持集上进行有效微调。它之所以“模型无关”是因为它不对模型架构 fθf**θ做强假设它可以应用于 CNN 或 RNN 等多种模型。其核心思想涉及一个两层优化过程内循环任务特定适应对于每个采样的任务 TiTi从当前元参数 θθ开始。仅使用任务的支持集 DisuppDisupp执行一次或几次梯度下降步骤以获得任务特定参数 θi′θ**i′。对于学习率为 αα的单个梯度步长θi′θ−α∇θLTi(fθ(Disupp))θ**i′θ−α∇θLTi(f**θ(Disupp))这里LTiLTi是任务 TiTi的损失函数fθ(Disupp)f**θ(Disupp) 表示模型使用参数 θθ对支持集进行预测的结果。请注意此梯度是相对于初始参数 θθ计算的。外循环元优化评估已适应参数 θi′θ**i′ 在任务查询集 DiqueryDiquery上的表现。元目标是在适应之后最小化跨任务的损失。元参数 θθ根据这些适应后查询集损失的总和或平均值进行更新使用元学习率 ββθ←θ−β∇θ∑Ti∼p(T)LTi(fθi′(Diquery))θ←θ−β∇θTi∼p(T)∑LTi(fθi′(Diquery))关键在于外循环中的梯度 ∇θ∑LTi(fθi′(…))∇θ∑LTi(fθi′(…)) 涉及到对内循环更新步骤的求导。这意味着我们需要计算相对于 θθ的梯度并考虑 θi′θ**i′ 是如何从 θθ推导出来的。这导致梯度计算涉及二阶导数梯度的梯度。内循环任务 Ti外循环跨任务元参数θ∇_θ L_supp(θ)计算支持集上的梯度已适应参数θ’适应步骤θ’ θ - α∇_θ∇_θ L_query(θ’)计算查询集上的梯度元更新(使用 ∇_θ)元梯度(涉及 ∇_θ’)更新 θ流程图说明了 MAML 优化过程。内循环使用支持集损失将参数 θθ适应为任务特定的 θ′θ′。外循环根据使用已适应参数 θ′θ′ 的查询集损失计算元梯度该元梯度随后用于更新原始元参数 θθ。在 PyTorch 中实现 MAML实现外循环的梯度计算需要谨慎。标准的 PyTorchbackward()调用会丢弃梯度中梯度计算所需的中间图信息。有两种主要的方法来处理这个问题使用torch.autograd.grad使用torch.autograd.grad并设置create_graphTrue参数来手动计算内部梯度。这会告诉 PyTorch 为梯度计算本身构建一个计算图从而允许稍后进行反向传播。示意图内循环梯度计算valinner_losscalculate_loss(model(support_x),support_y)valgradstorch.autograd.grad(inner_loss,model.parameters(),create_graphTrue)// 计算已适应参数函数式方法在这里通常更简单valadapted_params[p-alpha*gforp,g in zip(model.parameters(),grads)]// 使用 adapted_params 计算外部损失需要函数式模型调用// ... outer_loss calculate_loss(functional_model(adapted_params, query_x), query_y) ...// 外循环梯度计算稍后会汇总跨任务的 outer_loss// 并对总和调用 backward()。使用高阶梯度库像higher这样的库能显著简化此过程。higher提供了上下文管理器让你可以创建模型的临时可微分副本。你在此临时副本上执行内循环更新库会自动处理外循环梯度所需的计算跟踪。使用 ‘higher’ 的示意图importhighervalmeta_optimizer.zero_grad()valtotal_outer_loss0.0fortask_i-batch_of_tasks:val(support_x,support_y,query_x,query_y)get_task_data(task_i)withhigher.innerloop_ctx(model,inner_optimizer,copy_initial_weightsTrue)as(fmodel,diffopt):// 内循环更新for_-range(num_inner_steps):valinner_losscalculate_loss(fmodel(support_x),support_y)diffopt.step(inner_loss)// 更新 fmodel 的参数// 外循环评估valouter_losscalculate_loss(fmodel(query_x),query_y)total_outer_lossouter_loss// 反向传播元目标total_outer_loss.backward()meta_optimizer.step()higher方法因其更简洁的实现而常受青睐它抽象了create_graphTrue的手动处理和函数式参数更新。MAML 变体一阶 MAML (FOMAML)计算二阶导数在计算上可能开销很大。FOMAML 通过忽略二阶项来近似 MAML 更新。本质上它计算内部梯度 ∇θLTi(fθ(Disupp))∇θLTi(f**θ(Disupp))然后使用已适应参数计算外部梯度 ∇θ′LTi(fθi′(Diquery))∇θ′LTi(fθi′(Diquery))但在外部反向传播期间它将内部梯度步骤视为与初始 θθ无关。这更快但性能可能略逊于完整的 MAML。在 PyTorch 的手动实现中这对应于调用torch.autograd.grad不带create_graphTrue。Reptile另一种一阶元学习算法Nichol 等人2018它简化了更新过程。它在内循环中执行多个梯度步骤然后通过简单地将元参数 θθ稍微朝已适应参数 θi′θ**i′ 的方向移动来更新它们θ←θβ(θi′−θ)θ←θβ(θ**i′−θ)。这完全避免了显式的二阶导数计算。应用与考虑元学习特别是 MAML 及其变体已在以下方面获得应用少样本图像分类学习能够从极少量示例中识别新物体类别的分类器。强化学习训练能够快速适应新环境或动态变化的智能体。域适应将在一个数据分布源域上训练的模型适应到相关但不同分布目标域上并使用有限的目标数据实现良好性能。挑战计算成本MAML 的二阶梯度计算和存储开销可能很大特别是对于大型模型。FOMAML 和 Reptile 提供了替代方案。训练稳定性元学习优化场景可能很复杂有时需要仔细调整超参数例如内外部学习率、内循环步数。任务定义元学习的有效性很大程度上取决于任务 p(T)p(T) 的定义和分布。任务需要共享某种元学习器可以加以运用的潜在结构。元学习代表了一种转变从训练单个任务的模型转向训练具备高效学习能力的模型。像 MAML 这样的算法为实现这一目标提供了一个具体机制通过优化可适应的初始化使模型能够在数据稀缺的情况下快速适应。实现这些需要仔细处理梯度计算这通常通过专门的库或手动应用 PyTorch 的自动求导功能来简化。实践实现自定义GNN层PyTorch Geometric等库提供强大的预构建层来构建图神经网络GNN。虽然这些层非常有用但了解如何使用核心PyTorch操作从头开始构建GNN层能提供更全面的理解并具备实现新颖或定制化消息传递方案的灵活性。本次实践练习将指导你创建一个简单的自定义GNN层。许多GNN层背后的基本思想是消息传递即节点迭代地从邻居节点聚合信息并更新自身的表示。我们可以将其分解为每个节点 ii的两个主要步骤聚合从邻居节点 j∈N(i)j∈N(i) 收集特征或“消息”。更新将聚合的信息与节点当前的特征向量 hih**i结合以生成更新后的特征向量 hi′h**i′。让我们实现一个执行这些步骤的基本层。我们将定义一个层它使用可学习的权重矩阵转换节点特征使用简单的求和从邻居聚合转换后的特征然后应用激活函数。从数学上讲对于节点 ii此操作可以描述为ai∑j∈N(i)∪{i}Whja**ij∈N(i)∪{i}∑Whjhi′σ(ai)h**i′σ(a**i)这里hjh**j表示节点 jj的特征向量WW是一个可学习权重矩阵N(i)N(i) 是节点 ii的邻居集合σσ是一个非线性激活函数如ReLU。注意我们将节点自身ii也包含在聚合中这通常被称为添加自环。这确保了节点原始特征在更新时得到考量。设置自定义层首先请确保已导入PyTorch。我们将把自定义层定义为一个继承自torch.nn.Module的Python类。importtorchimporttorch.nnas nnimporttorch.nn.functionalas FclassSimpleGNNLayerextendsnn.Module: 一个实现消息传递的基本图神经网络层。 Args: in_features (int): 每个输入节点特征向量的大小。 out_features (int): 每个输出节点特征向量的大小。 def__init__(in_features:Int,out_features:Int):super(SimpleGNNLayer,self).__init__()valin_featuresin_featuresvalout_featuresout_features// 定义可学习的权重矩阵vallinearnn.Linear(in_features,out_features,biasFalse)// 初始化权重可选但通常是好的做法nn.init.xavier_uniform_(linear.weight)defforward(x:torch.Tensor,edge_index:torch.Tensor): 定义每次调用时执行的计算。 Args: x (torch.Tensor): 节点特征张量形状为 [num_nodes, in_features]。 edge_index (torch.Tensor): COO格式的图连接信息形状为 [2, num_edges]。 edge_index[0] 源节点edge_index[1] 目标节点。 Returns: torch.Tensor: 更新后的节点特征张量形状为 [num_nodes, out_features]。 valnum_nodesx.size(0)// 1. 为edge_index表示的邻接矩阵添加自环// 创建节点索引张量 [0, 1, ..., num_nodes-1]valself_loopstorch.arange(0,num_nodes,devicex.device).unsqueeze(0)valself_loopsself_loops.repeat(2,1)// 形状 [2, num_nodes]// 将原始边与自环拼接valedge_index_with_self_loopstorch.cat([edge_index,self_loops],dim1)// 提取源节点和目标节点索引valrowedge_index_with_self_loops(0)valcoledge_index_with_self_loops(1)// 2. 线性变换节点特征valx_transformedlinear(x)// 形状: [num_nodes, out_features]// 3. 聚合来自邻居包括自身的特征// 我们希望对每个目标节点col求和源节点row的特征// 使用零初始化输出张量valaggregated_featurestorch.zeros(num_nodes,out_features,devicex.device)// 使用 index_add_ 进行高效聚合散列求和// 将 x_transformed[row] 的元素添加到 aggregated_features 中由 col 指定的索引处// index_add_(维度, 索引张量, 要添加的张量)aggregated_features.index_add_(0,col,x_transformed(row))// 4. 应用最终激活函数可选// 在此示例中我们使用ReLUvaloutput_featuresF.relu(aggregated_features)returnoutput_featuresdef__repr__(self):returnf{self.__class__.__name__}({self.in_features}, {self.out_features})理解实现初始化 (__init__)我们定义一个nn.Linear层。此层将把可学习权重变换 WW应用于输入节点特征。为简单起见我们设置biasFalse这与一些GNN公式如基本GCN一致。使用nn.init.xavier_uniform_进行权重初始化有助于稳定训练。前向传播 (forward)这是消息传递逻辑的所在。自环我们显式地将自环添加到edge_index。这确保了在为节点聚合邻居特征时节点自身的转换特征也包含在内。我们创建一个表示从每个节点到自身的边的边索引并将其与原始edge_index拼接。特征变换我们同时对所有节点特征x应用线性变换 (self.linear)。聚合这是GNN的主要步骤。我们需要对每个目标节点 (col) 求和源节点 (x_transformed[row]) 的转换特征。torch.index_add_是一种高效执行此“散列-求和”操作的方法。它接受要累积到的张量 (aggregated_features)、进行索引的维度节点为0、要添加到的索引 (col即目标节点以及要添加的值 (x_transformed[row]即源节点的转换特征。激活最后逐元素应用一个非线性激活函数 (F.relu)。这里有一个小型图可视化用以显示edge_index格式和邻居的思想0123对于上面的图一个可能的edge_index表示用于消息传递的有向边假设无向原始边意味着消息双向传递可能是tensor([[0, 0, 1, 2, 1, 2, 3, 3], [1, 2, 0, 0, 3, 3, 1, 2]])。 第一行包含源节点第二行包含目标节点。当为节点3聚合时我们会查看来自源节点1和2的消息。使用自定义层现在让我们看看如何使用这个SimpleGNNLayer。我们需要一些示例节点特征和一个edge_index。// 示例用法// 定义图数据valnum_nodes4valnum_features8valout_layer_features16// 节点特征随机valxtorch.randn(num_nodes,num_features)// 边索引表示连接例如0-1, 0-2, 1-3, 2-3对于无向图则反之valedge_indextorch.tensor(Seq(Seq(0,0,1,2,1,2,3,3),// 源节点Seq(1,2,0,0,3,3,1,2)// 目标节点),dtypetorch.long)// 实例化层valgnn_layerSimpleGNNLayer(in_featuresnum_features,out_featuresout_layer_features)println(s已实例化层:$gnn_layer)// 将数据通过该层valoutput_node_featuresgnn_layer(x,edge_index)// 检查输出形状println(s\n输入节点特征形状:${x.shape})println(s边索引形状:${edge_index.shape})println(s输出节点特征形状:${output_node_features.shape})// 验证输出形状是否符合预期: [num_nodes, out_features]assert output_node_features.shape(num_nodes,out_layer_features)print(\n数据已成功通过自定义GNN层。)// 显示节点0的前几个输出特征println(s节点0的输出特征前5维:${output_node_features(0,0until5)})此示例展示了创建随机节点特征和示例edge_index实例化我们的SimpleGNNLayer并执行前向传播。输出形状[num_nodes, out_features]确认该层按预期运行为每个节点根据其邻域生成新的嵌入。潜在的扩展这个简单的层可作为根本。你可以通过多种方式对其进行扩展不同聚合方式将index_add_求和聚合替换为平均或最大值聚合。平均聚合通常需要知道每个节点的度。边特征修改forward传播以接受和运用边特征并可能在聚合前将其加入到消息计算中。标准化添加标准化步骤例如GCN层中常见的对称标准化这通常涉及节点度。偏置项在nn.Linear层中包含一个偏置项或在聚合后添加。多层堆叠堆叠这些层可能加入标准化或跳跃连接以构建更深的GNN模型。构建这样的自定义层是一项很有价值的技能。它使你能够直接根据研究论文实现前沿GNN架构或在必要时精确地根据问题需求定制消息传递方案。构建自定义nn.Module组件的这一相同原理也适用于在本课程中实现的Transformer、归一化流或其他高级架构中的独特机制。