Package: gradient-descent-the-ultimate-optimizer
Gradient Descent: The Ultimate OptimizerAbstractWorking with any gradient-based machine learning algorithm involves the tedious task of tuning the optimizer's hyperparameters, such as the step size. Recent work has shown how the step size can itself be "learned" on-line by gradient descent, by manually deriving expressions for "hypergradients" ahead of time.We show how toautomaticallycompute hypergradients with a simple and elegant modification to backpropagation. This allows us to apply the method to other hyperparameters besides the step size, such as the momentum coefficient. We can even recursively apply the method to its ownhyper-hyperparameters, and so onad infinitum. As these towers of optimizers grow taller, they become less sensitive to the initial choice of hyperparameters. We present experiments validating this for MLPs, CNNs, and RNNs.This repository contains an implementation of the algorithm in our paper.Citation@article{chandra2022gradient,
    title = {Gradient Descent: The Ultimate Optimizer},
    author = {Chandra, Kartik and Xie, Audrey and Ragan-Kelley, Jonathan and Meijer, Erik},
    journal = {NeurIPS},
    year = {2022},
    url = {https://arxiv.org/abs/1909.13371}
}Install# install pytorch for your specific machine
...

# install our package
pip install gradient-descent-the-ultimate-optimizerExampleFirst, build the MLP and initialize data loaders as you would normally in PyTorch.importmathimporttorchimporttorchvisionimporttorch.nnasnnimporttorch.nn.functionalasFclassMNIST_FullyConnected(nn.Module):"""A fully-connected NN for the MNIST task. This is Optimizable but not itselfan optimizer."""def__init__(self,num_inp,num_hid,num_out):super(MNIST_FullyConnected,self).__init__()self.layer1=nn.Linear(num_inp,num_hid)self.layer2=nn.Linear(num_hid,num_out)definitialize(self):nn.init.kaiming_uniform_(self.layer1.weight,a=math.sqrt(5))nn.init.kaiming_uniform_(self.layer2.weight,a=math.sqrt(5))defforward(self,x):"""Compute a prediction."""x=self.layer1(x)x=torch.tanh(x)x=self.layer2(x)x=torch.tanh(x)x=F.log_softmax(x,dim=1)returnxBATCH_SIZE=256EPOCHS=5DEVICE='cuda'iftorch.cuda.is_available()else'cpu'mnist_train=torchvision.datasets.MNIST('./data',train=True,download=True,transform=torchvision.transforms.ToTensor())mnist_test=torchvision.datasets.MNIST('./data',train=False,download=True,transform=torchvision.transforms.ToTensor())dl_train=torch.utils.data.DataLoader(mnist_train,batch_size=BATCH_SIZE,shuffle=True)dl_test=torch.utils.data.DataLoader(mnist_test,batch_size=10000,shuffle=False)model=MNIST_FullyConnected(28*28,128,10).to(DEVICE)Next, import our package and initialize a stack of hyperoptimizers. This example uses the stackAdam/SGD.fromgradient_descent_the_ultimate_optimizerimportgdtuooptim=gdtuo.Adam(optimizer=gdtuo.SGD(1e-5))gdtuo.ModuleWrapperallows anynn.Moduleto be optimized by hyperoptimizers.mw=gdtuo.ModuleWrapper(model,optimizer=optim)mw.initialize()Lastly, usemwinstead of a PyTorch optimizer to optimize the model. The train loop is nearly identical to what you would typically implement in PyTorch (differences are marked by comments).foriinrange(1,EPOCHS+1):running_loss=0.0forj,(features_,labels_)inenumerate(dl_train):mw.begin()# call this before each step, enables gradient tracking on desired paramsfeatures,labels=torch.reshape(features_,(-1,28*28)).to(DEVICE),labels_.to(DEVICE)pred=mw.forward(features)loss=F.nll_loss(pred,labels)mw.zero_grad()loss.backward(create_graph=True)# important! use create_graph=Truemw.step()running_loss+=loss.item()*features_.size(0)train_loss=running_loss/len(dl_train.dataset)print("EPOCH:{}, TRAIN LOSS:{}".format(i,train_loss))Note that on the first step of the train loop PyTorch will return the following warning:UserWarning: Using backward() with create_graph=True will create a reference cycle between the parameter and its gradient which can cause a memory leak. We recommend using autograd.grad when creating the graph to avoid this. If you have to use this function, make sure to reset the .grad fields of your parameters to None after use to break the cycle and avoid the leak.This is normal and to be expected.