Package: torchopt
Installation|Documentation|Tutorials|Examples|Paper|CitationTorchOptis an efficient library for differentiable optimization built uponPyTorch.
TorchOpt is:Comprehensive: TorchOpt provides three differentiation modes - explicit differentiation, implicit differentiation, and zero-order differentiation for handling different differentiable optimization situations.Flexible: TorchOpt provides both functional and objective-oriented API for users' different preferences. Users can implement differentiable optimization in JAX-like or PyTorch-like style.Efficient: TorchOpt provides (1) CPU/GPU acceleration differentiable optimizer (2) RPC-based distributed training framework (3) Fast Tree Operations, to largely increase the training efficiency for bi-level optimization problems.Beyond differentiable optimization, TorchOpt can also be regarded as a functional optimizer that enablesJAX-likecomposable functional optimizer for PyTorch.
With TorchOpt, users can easily conduct neural network optimization in PyTorch with a functional style optimizer, similar toOptaxin JAX.The README is organized as follows:TorchOpt as Functional OptimizerOptax-Like APIPyTorch-Like APIDifferentiableTorchOpt for Differentiable OptimizationExplicit Gradient (EG)Implicit Gradient (IG)Zero-order Differentiation (ZD)High-Performance and Distributed TrainingCPU/GPU accelerated differentiable optimizerDistributed TrainingOpTreeVisualizationExamplesInstallationChangelogCiting TorchOptThe TeamLicenseTorchOpt as Functional OptimizerThe design of TorchOpt follows the philosophy of functional programming.
Aligned withfunctorch, users can conduct functional style programming with models, optimizers and training in PyTorch.
We use the Adam optimizer as an example in the following illustration.
You can also check out the tutorial notebookFunctional Optimizerfor more details.Optax-Like APIFor those users who prefer fully functional programming, we offer Optax-Like API by passing gradients and optimizer states to the optimizer function.
Here is an example coupled withfunctorch:classNet(nn.Module):...classLoader(DataLoader):...net=Net()# initloader=Loader()optimizer=torchopt.adam()model,params=functorch.make_functional(net)# use functorch extract network parametersopt_state=optimizer.init(params)# init optimizerxs,ys=next(loader)# get datapred=model(params,xs)# forwardloss=F.cross_entropy(pred,ys)# compute lossgrads=torch.autograd.grad(loss,params)# compute gradientsupdates,opt_state=optimizer.update(grads,opt_state)# get updatesparams=torchopt.apply_updates(params,updates)# update network parametersWe also provide a wrappertorchopt.FuncOptimizerto make maintaining the optimizer state easier:net=Net()# initloader=Loader()optimizer=torchopt.FuncOptimizer(torchopt.adam())# wrap with `torchopt.FuncOptimizer`model,params=functorch.make_functional(net)# use functorch extract network parametersforxs,ysinloader:# get datapred=model(params,xs)# forwardloss=F.cross_entropy(pred,ys)# compute lossparams=optimizer.step(loss,params)# update network parametersPyTorch-Like APIWe also design a base classtorchopt.Optimizerthat has the same interface astorch.optim.Optimizer.
We offer origin PyTorch APIs (e.g.zero_grad()orstep()) by wrapping our Optax-Like API for traditional PyTorch users.net=Net()# initloader=Loader()optimizer=torchopt.Adam(net.parameters())xs,ys=next(loader)# get datapred=net(xs)# forwardloss=F.cross_entropy(pred,ys)# compute lossoptimizer.zero_grad()# zero gradientsloss.backward()# backwardoptimizer.step()# step updatesDifferentiableOn top of the same optimization function astorch.optim, an important benefit of the functional optimizer is that one can implement differentiable optimization easily.
This is particularly helpful when the algorithm requires differentiation through optimization updates (such as meta-learning practices).
We take as the inputs the gradients and optimizer states, and use non-in-place operators to compute and output the updates.
The processes can be automatically implemented, with the only need from users being to pass the argumentinplace=Falseto the functions.
Check out the sectionExplicit Gradient (EG)functional API for example.TorchOpt for Differentiable OptimizationWe design a bilevel-optimization updating scheme, which can be easily extended to realize various differentiable optimization processes.As shown above, the scheme contains an outer level that has parameters $\phi$ that can be learned end-to-end through the inner level parameters solution $\theta^{\prime}(\phi)$ by using the best-response derivatives $\partial \theta^{\prime}(\phi) / \partial \phi$.
TorchOpt supports three differentiation modes.
It can be seen that the key component of this algorithm is to calculate the best-response (BR) Jacobian.
From the BR-based perspective, existing gradient methods can be categorized into three groups: explicit gradient over unrolled optimization, implicit differentiation, and zero-order gradient differentiation.Explicit Gradient (EG)The idea of the explicit gradient is to treat the gradient step as a differentiable function and try to backpropagate through the unrolled optimization path.
This differentiation mode is suitable for algorithms when the inner-level optimization solution is obtained by a few gradient steps, such asMAMLandMGRL.
TorchOpt offers both functional and object-oriented API for EG to fit different user applications.Functional APIThe functional API is to conduct optimization in a functional programming style.
Note that we pass the argumentinplace=Falseto the functions to make the optimization differentiable.
Refer to the tutorial notebookFunctional Optimizerfor more guidance.# Define functional optimizeroptimizer=torchopt.adam()# Define meta and inner parametersmeta_params=...fmodel,params=make_functional(model)# Initial statestate=optimizer.init(params)foriterinrange(iter_times):loss=inner_loss(fmodel,params,meta_params)grads=torch.autograd.grad(loss,params)# Apply non-inplace parameter updateupdates,state=optimizer.update(grads,state,inplace=False)params=torchopt.apply_updates(params,updates)loss=outer_loss(fmodel,params,meta_params)meta_grads=torch.autograd.grad(loss,meta_params)OOP APITorchOpt also provides OOP API compatible with the PyTorch programming style.
Refer to the example and the tutorial notebookMeta-Optimizer,Stop Gradientfor more guidance.# Define meta and inner parametersmeta_params=...model=...# Define differentiable optimizeroptimizer=torchopt.MetaAdam(model)# a model instance as the argument instead of model.parameters()foriterinrange(iter_times):# Perform inner updateloss=inner_loss(model,meta_params)optimizer.step(loss)loss=outer_loss(model,meta_params)loss.backward()Implicit Gradient (IG)By treating the solution $\theta^{\prime}$ as an implicit function of $\phi$, the idea of IG is to directly get analytical best-response derivatives $\partial \theta^{\prime} (\phi) / \partial \phi$ byimplicit function theorem.
This is suitable for algorithms when the inner-level optimal solution is achieved ${\left. \frac{\partial F (\theta, \phi)}{\partial \theta} \right\rvert}_{\theta=\theta^{\prime}} = 0$ or reaches some stationary conditions $F (\theta^{\prime}, \phi) = 0$, such asiMAMLandDEQ.
TorchOpt offers both functional and OOP APIs for supporting bothconjugate gradient-basedandNeumann series-basedIG methods.
Refer to the exampleiMAMLand the notebookImplicit Gradientfor more guidance.Functional APIFor the implicit gradient, similar toJAXopt, users need to define the stationary condition and TorchOpt provides the decorator to wrap the solve function for enabling implicit gradient computation.# The stationary condition for the inner-loopdefstationary(params,meta_params,data):# Stationary condition constructionreturnstationarycondition# Decorator for wrapping the function# Optionally specify the linear solver (conjugate gradient or Neumann series)@torchopt.diff.implicit.custom_root(stationary,solve=linear_solver)defsolve(params,meta_params,data):# Forward optimization process for paramsreturnoutput# Define params, meta_params and get dataparams,meta_prams,data=...,...,...optimal_params=solve(params,meta_params,data)loss=outer_loss(optimal_params)meta_grads=torch.autograd.grad(loss,meta_params)OOP APITorchOpt also offers an OOP API, which users need to inherit from the classtorchopt.nn.ImplicitMetaGradientModuleto construct the inner-loop network.
Users need to define the stationary condition/objective function and the inner-loop solve function to enable implicit gradient computation.# Inherited from the class ImplicitMetaGradientModule# Optionally specify the linear solver (conjugate gradient or Neumann series)classInnerNet(ImplicitMetaGradientModule,linear_solve=linear_solver):def__init__(self,meta_param):super().__init__()self.meta_param=meta_param...defforward(self,batch):# Forward process...defoptimality(self,batch,labels):# Stationary condition construction for calculating implicit gradient# NOTE: If this method is not implemented, it will be automatically# derived from the gradient of the `objective` function....defobjective(self,batch,labels):# Define the inner-loop optimization objective...defsolve(self,batch,labels):# Conduct the inner-loop optimization...# Get meta_params and datameta_params,data=...,...inner_net=InnerNet(meta_params)# Solve for inner-loop process related to the meta-parametersoptimal_inner_net=inner_net.solve(data)# Get outer loss and solve for meta-gradientloss=outer_loss(optimal_inner_net)meta_grads=torch.autograd.grad(loss,meta_params)Zero-order Differentiation (ZD)When the inner-loop process is non-differentiable or one wants to eliminate the heavy computation burdens in the previous two modes (brought by Hessian), one can choose Zero-order Differentiation (ZD).
ZD typically gets gradients based on zero-order estimation, such as finite-difference, orEvolutionary Strategy.
Instead of optimizing the objective $F$, ES optimizes a smoothed objective.
TorchOpt provides both functional and OOP APIs for the ES method.
Refer to the tutorial notebookZero-order Differentiationfor more guidance.Functional APIFor zero-order differentiation, users need to define the forward pass calculation and the noise sampling procedure. TorchOpt provides the decorator to wrap the forward function for enabling zero-order differentiation.# Customize the noise sampling function in ESdefdistribution(sample_shape):# Generate a batch of noise samples# NOTE: The distribution should be spherical symmetric and with a constant variance of 1....returnnoise_batch# Distribution can also be an instance of `torch.distributions.Distribution`, e.g., `torch.distributions.Normal(...)`distribution=torch.distributions.Normal(loc=0,scale=1)# Specify method and hyper-parameter of ES@torchopt.diff.zero_order(distribution,method)defforward(params,batch,labels):# Forward process...returnobjective# the returned tensor should be a scalar tensorOOP APITorchOpt also offers an OOP API, which users need to inherit from the classtorchopt.nn.ZeroOrderGradientModuleto construct the network as annn.Modulefollowing a classical PyTorch style.
Users need to define the forward process zero-order gradient proceduresforward()and a noise sampling functionsample().# Inherited from the class ZeroOrderGradientModule# Optionally specify the `method` and/or `num_samples` and/or `sigma` used for samplingclassNet(ZeroOrderGradientModule,method=method,num_samples=num_samples,sigma=sigma):def__init__(self,...):...defforward(self,batch):# Forward process...returnobjective# the returned tensor should be a scalar tensordefsample(self,sample_shape=torch.Size()):# Generate a batch of noise samples# NOTE: The distribution should be spherical symmetric and with a constant variance of 1....returnnoise_batch# Get model and datanet=Net(...)data=...# Forward passloss=Net(data)# Backward pass using zero-order differentiationgrads=torch.autograd.grad(loss,net.parameters())High-Performance and Distributed TrainingCPU/GPU accelerated differentiable optimizerWe take the optimizer as a whole instead of separating it into several basic operators (e.g.,sqrtanddiv).
Therefore, by manually writing the forward and backward functions, we can perform the symbolic reduction.
In addition, we can store some intermediate data that can be reused during the backpropagation.
We write the accelerated functions in C++ OpenMP and CUDA, bind them bypybind11to allow they can be called by Python, and then define the forward and backward behavior usingtorch.autograd.Function.
Users can use it by simply setting theuse_accelerated_opflag asTrue.
Refer to the corresponding sections in the tutorialsFunctional Optimizer](tutorials/1_Functional_Optimizer.ipynb) andMeta-Optimizeroptimizer=torchopt.MetaAdam(model,lr,use_accelerated_op=True)Distributed TrainingTorchOptprovides distributed training features based on the PyTorch RPC module for better training speed and multi-node multi-GPU support.
Different from the MPI-like parallelization paradigm, which uses multiple homogeneous workers and requires carefully designed communication hooks, the RPC APIs allow users to build their optimization pipeline more flexibly.
Experimental results show that we achieve an approximately linear relationship between the speed-up ratio and the number of workers.
Check out theDistributed Training Documentationanddistributed MAML examplefor more specific guidance.OpTreeWe implement thePyTreeto enable fast nested structure flattening using C++.
The tree operations (e.g., flatten and unflatten) are very important in enabling functional and Just-In-Time (JIT) features of deep learning frameworks.
By implementing it in C++, we can use some cache/memory-friendly structures (e.g.,absl::InlinedVector) to improve the performance.
For more guidance and comparison results, please refer to our open-source projectOpTree.VisualizationComplex gradient flow in meta-learning brings in a great challenge for managing the gradient flow and verifying its correctness of it.
TorchOpt provides a visualization tool that draws variable (e.g., network parameters or meta-parameters) names on the gradient graph for better analysis.
The visualization tool is modified fromtorchviz.
Refer to the examplevisualization codeand the tutorial notebookVisualizationfor more details.The figure below shows the visualization result.
Compared withtorchviz, TorchOpt fuses the operations within theAdamtogether (orange) to reduce the complexity and provide simpler visualization.ExamplesIn theexamplesdirectory, we offer several examples of functional optimizers and lightweight meta-learning examples with TorchOpt.Model-Agnostic Meta-Learning (MAML) - Supervised Learning(ICML 2017)Learning to Reweight Examples for Robust Deep Learning(ICML 2018)Model-Agnostic Meta-Learning (MAML) - Reinforcement Learning(ICML 2017)Meta-Gradient Reinforcement Learning (MGRL)(NeurIPS 2018)Learning through opponent learning process (LOLA)(AAMAS 2018)Meta-Learning with Implicit Gradients(NeurIPS 2019)Also, checkexamplesfor more distributed/visualization/functorch-compatible examples.InstallationRequirementsPyTorch(Optional) For visualizing computation graphsGraphviz(for Linux users useapt/yum install graphvizorconda install -c anaconda python-graphviz)Please follow the instructions athttps://pytorch.orgto install PyTorch in your Python environment first.Then run the following command to install TorchOpt from PyPI (/):pip3installtorchoptIf the minimum version of PyTorch is not satisfied,pipwill install/upgrade it for you. Please be careful about thetorchbuild for CPU / CUDA support (e.g.cpu,cu116,cu117).
You may need to specify the extra index URL for thetorchpackage:pip3installtorchopt--extra-index-urlhttps://download.pytorch.org/whl/cu117Seehttps://pytorch.orgfor more information about installing PyTorch.You can also build shared libraries from source, use:gitclonehttps://github.com/metaopt/torchopt.gitcdtorchopt
pip3install.We provide acondaenvironment recipe to install the build toolchain such ascmake,g++, andnvcc.
You can use the following commands withconda/mambato create a new isolated environment.gitclonehttps://github.com/metaopt/torchopt.gitcdtorchopt# You may need `CONDA_OVERRIDE_CUDA` if conda fails to detect the NVIDIA driver (e.g. in docker or WSL2)CONDA_OVERRIDE_CUDA=11.7condaenvcreate--fileconda-recipe-minimal.yaml

condaactivatetorchopt
makeinstall-editable# or run `pip3 install --no-build-isolation --editable .`ChangelogSeeCHANGELOG.md.Citing TorchOptIf you find TorchOpt useful, please cite it in your publications.@article{torchopt,title={TorchOpt: An Efficient Library for Differentiable Optimization},author={Ren, Jie and Feng, Xidong and Liu, Bo and Pan, Xuehai and Fu, Yao and Mai, Luo and Yang, Yaodong},journal={arXiv preprint arXiv:2211.06934},year={2022}}The TeamTorchOpt is a work byJie Ren,Xidong Feng,Bo Liu,Xuehai Pan,Luo Mai, andYaodong Yang.LicenseTorchOpt is released under the Apache License, Version 2.0.