# Project export: Loss Landscapes in a Lens This document was generated by HackStack to give an AI agent context about a hackathon project. Sections are labeled with their provenance; content marked as truncated was cut to keep this document small. ## Project metadata - Hackathon: Cal Hacks 11.0 - Tagline: Rendering neural net loss landscapes in real life using Snap Spectacles - Devpost: https://devpost.com/software/loss-landscapes - GitHub: https://github.com/amosyou/loss-landscapes-lens - Team: 1 GitHub contributor(s) — Amos You (9 commits) ## Devpost submission (written by the team) ### Inspiration This project was inspired by Visualizing the Loss Landscape of Neural Nets, a NeurIPS 2018 paper. The paper describes a method for visualizing loss landscapes, which are graphs that reveal the error of neural network predictions as a function of their parameters. These visualizations are useful for understanding how the neural network is learning during the training process, and provides insight into the generalization of these models. These 3D surface plots are really cool to view on plotting libraries / 3D modeling software, but I've always wished I could move around and view the shape of the meshes in real life. ### What it does This project directly applies the Visualizing the Loss Landscape of Neural Nets paper by rendering these loss landscapes using the Snap Spectacles. With wearable tech and smart glasses like the Snap Spectacles, we can finally have an immersive experience in AR to better visualize the shapes of these 3D surface plots, and have visibility into more fine-grain or precise details. The differences that the paper finds when comparing how different model architectures lead to different loss surfaces and optimization trajectories can be clearly seen when viewing through the Snap Spectacles. ### How we built it The loss-landscape repo provides code to generate .vtp files for the 3D loss surface plots of a few neural net models. Lens Studio only allows for 3D object imports from .fbx, .obj, and .gltf file formats, so we convert the .vtp files to .obj files using PyVista, a 3D plotting Python library. Then, we import these loss surface meshes into Lens Studio as assets, and customize the arrangement of the camera, change scaling and positioning of the mesh, enable surface detection, and apply material to the mesh for the color gradient. And with some sync'ing to the Spectacles, we have a working lens! ### Challenges we ran into There were many challenges on the AR development side, as this was our first time working with AR. One notable challenge was shading. Shading the mesh was not an easy task, but we found a hacky solution that uses the Twist material, disabling its animation, and warping the mesh slightly, which gives our desired color gradient effect. We also began fine-tuning some LLMs (ie. Llama 3), but didn't have enough time to generate the 3D surface plots for the loss landscape. ### Accomplishments we're proud of Considering that this was our first time building with AR, we felt immediate wins even when we created our first barebones version of our lens. In only a few hours, we were able to get our Snap Spectacles set up and a lens with our desired outcome. ### What we learned Developing the Spectacles on Lens Studio was very pleasant. The software provides an UI similar to Premiere Pro where you import and configure assets on the left, with your simulated scene on the right, making the development process quite intuitive for us (and probably other content creators as well). We learned a lot on the different types of 3D modeling software out there and gained exposure to how shaders work with a hint of OpenGL. ### What's next for Loss Landscapes Although we have working code to generate the 3D surface plots given some model, we're hoping there's a way to automate the setup of the assets in Lens Studio. We're also hoping to continue fine-tuning some LLMs such as Llama 3 and diffusion models like FLUX.1-dev to see how the loss surfaces may have changed for newer models compared to the traditional models (ie. DenseNet, ResNet, and VGG) presented in the paper. Acknowledgements [1] thumbnail modified from here [2] loss-landscape repo ## README (from the GitHub repository) # Loss Landscapes in a Lens ![loss landscape on spectacles](assets/lens.gif) This project was inspired by [Visualizing the Loss Landscape of Neural Nets](https://arxiv.org/abs/1712.09913), a NeurIPS 2018 paper. The paper describes a method for visualizing loss landscapes of neural networks by applying normalization to specific model parameters, and then plotting loss values along 2 randomly selected directions during training. This project directly applies the paper by rendering these loss landscapes in real life using the Snap Spectacles. With wearable tech and smart glasses like the Snap Spectacles, we can finally have an immersive experience in AR to better visualize the shapes of these 3D surface plots, and have visibility into more fine-grain or precise details. The lens can be found [here](https://www.snapchat.com/lens/724f44537ca84400acef41af5ae386b3). ## Usage 1. To install dependencies, create a Python virtual environment and run the following command. ```bash pip install -r requirements.txt ``` 2. `loss-landscape/` corresponds to the [loss-landscape](https://github.com/tomgoldstein/loss-landscape) repo, which provides code to generate .vtp files for the 3D loss surface plots of a few models. 3. Lens Studio only allows for 3D object imports from .fbx, .obj, and .gltf file formats, so we convert the .vtp files to .obj files using PyVista, a 3D plotting Python library. The notebook can be found at `vtp2obj.ipynb`. 4. `lens/` holds the Lens Studio project with the project file and assets. We import these loss surface meshes into Lens Studio as assets, and customize the arrangement of the camera, change scaling and positioning of the mesh, enable surface detection, and apply material to the mesh for the color gradient. And with some sync'ing to the Spectacles, we have a working lens! ## Future Work Although we have working code to generate the 3D surface plots given some model, we're hoping there's a way to automate the setup of the assets in Lens Studio. We're also hoping to continue fine-tuning some LLMs such as Llama 3 and diffusion models like FLUX.1-dev to see how the loss surfaces differ for newer models compared to the traditional models (ie. DenseNet, ResNet, and VGG) presented in the paper. ## Acknowledgements Code heavily borrowed from [loss-landscape](https://github.com/tomgoldstein/loss-landscape). ## Detected evidence (automated analysis) Indexed codebase: 41 recognized source files, 153 KB. - Python (language) — detected in the code - PyTorch (technology) — detected in the code ## Codebase structure (from repository index) ### Files (81 of 81) ``` .gitignore lens/Assets/.DS_Store lens/Assets/Base.png.meta lens/Assets/Device Camera Texture.deviceCameraTexture lens/Assets/Device Camera Texture.deviceCameraTexture.meta lens/Assets/Material.mat lens/Assets/Material.mat.meta lens/Assets/MaterialParams.png.meta lens/Assets/Normal.png.meta lens/Assets/PBR.mat lens/Assets/PBR.mat.meta lens/Assets/pbr.ss_graph lens/Assets/pbr.ss_graph.meta lens/Assets/Render Target.renderTarget lens/Assets/Render Target.renderTarget.meta lens/Assets/resnet.mtl lens/Assets/resnet.mtl.meta lens/Assets/resnet.obj lens/Assets/resnet.obj.meta lens/Assets/Scene.scene lens/Assets/Scene.scene.meta lens/Assets/ScreenTexture.screenTexture lens/Assets/ScreenTexture.screenTexture.meta lens/Assets/Surface Detection.lspkg lens/Assets/Surface Detection.lspkg.meta lens/Assets/Twist Material/.hidden/Twist.ss_graph lens/Assets/Twist Material/.hidden/Twist.ss_graph.meta lens/Assets/Twist Material/Twist.mat lens/Assets/Twist Material/Twist.mat.meta lens/Assets/Uber PBR.mat lens/Assets/Uber PBR.mat.meta lens/Assets/uber_pbr.ss_graph lens/Assets/uber_pbr.ss_graph.meta lens/loss.esproj lens/PluginsUserPreferences/Snap_Plugin_Core_AssetBrowserService.yaml lens/PluginsUserPreferences/Snap_Plugin_Gui_AssetBrowser_cd11ace2-8507-4965-a068-d85c26031a6b.yaml lens/PluginsUserPreferences/Snap_Plugin_Gui_Graph_Editor_70837bf1-0fe0-4d93-b6c8-5d6036a3da9c.yaml lens/PluginsUserPreferences/Snap_Plugin_Gui_InspectorPanel_780d9164-e681-4de1-b586-5213ed5523b9.yaml lens/PluginsUserPreferences/Snap_Plugin_Gui_InspectorPanel_9cebccd3-697a-425f-a96d-16a1622b0ab3.yaml lens/PluginsUserPreferences/Snap_Plugin_Gui_LoggerPanel_81de965d-0c94-4d28-82cf-901740003a1f.yaml lens/PluginsUserPreferences/Snap_Plugin_Gui_PreviewPanel_52fc2a95-41a3-427b-a138-1e00eb1c405e.yaml lens/PluginsUserPreferences/Snap_Plugin_Gui_SceneEditor_a14d85d7-9a15-42a6-913e-3f24dea267c2.yaml lens/PluginsUserPreferences/Snap_Plugin_Gui_SceneHierarchyEditor_6778142c-0629-4c2a-afc3-92a65ce29e1a.yaml loss-landscape/.gitignore loss-landscape/cifar10/__init__.py loss-landscape/cifar10/dataloader.py loss-landscape/cifar10/main.py loss-landscape/cifar10/model_loader.py loss-landscape/cifar10/models/__init__.py loss-landscape/cifar10/models/densenet.py loss-landscape/cifar10/models/resnet.py loss-landscape/cifar10/models/vgg.py loss-landscape/dataloader.py loss-landscape/evaluation.py loss-landscape/h5_util.py loss-landscape/h52vtp.py loss-landscape/hess_vec_prod.py loss-landscape/LICENSE loss-landscape/model_loader.py loss-landscape/mpi4pytorch.py loss-landscape/net_plotter.py loss-landscape/plot_1D.py loss-landscape/plot_2D.py loss-landscape/plot_hessian_eigen.py loss-landscape/plot_surface.py loss-landscape/plot_trajectory.py loss-landscape/projection.py loss-landscape/README.md loss-landscape/scheduler.py loss-landscape/script/1d_linear_interpolation_resnet56.sh loss-landscape/script/1d_linear_interpolation_vgg9.sh loss-landscape/script/1d_loss_curve_resnet56.sh loss-landscape/script/1d_loss_curve_vgg9.sh loss-landscape/script/2d_contour_resnet_noshort.sh loss-landscape/script/2d_contour_resnet56.sh loss-landscape/script/2d_contour_vgg9.sh README.md requirements.txt resnet56.mtl resnet56.obj vtp2obj.ipynb ``` ### Dependencies - requirements.txt: h5py, jupyter, matplotlib, mpi4py, numpy, openmpi, pyvista, scikit-learn, scipy, seaborn, torch, torchvision ### Recent commits (newest first) - update README - feat: add lens studio project - feat: add gif teaser - update README - update README - feat: convert vtp to obj file using pyvista - feat: add requirements from loss landscape - feat: add loss landscape repo - Initial commit ## Key source files (fetched from GitHub, selected and truncated for size) ### requirements.txt ``` torch openmpi mpi4py numpy h5py matplotlib scipy torchvision scikit-learn seaborn pyvista jupyter ``` ### loss-landscape/cifar10/main.py ```python from __future__ import print_function import os import random import numpy as np import argparse import torch import torch.nn as nn import torch.nn.init as init import torch.optim as optim import torch.nn.functional as F import torch.backends.cudnn as cudnn from torch.autograd import Variable import torch.nn.parallel import model_loader import dataloader def init_params(net): for m in net.modules(): if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d): init.kaiming_normal_(m.weight, mode='fan_in') if m.bias is not None: init.constant_(m.bias, 0) elif isinstance(m, nn.BatchNorm2d): init.constant_(m.weight, 1) init.constant_(m.bias, 0) elif isinstance(m, nn.Linear): init.normal_(m.weight, std=1e-3) if m.bias is not None: init.constant_(m.bias, 0) # Training def train(trainloader, net, criterion, optimizer, use_cuda=True): net.train() train_loss = 0 correct = 0 total = 0 if isinstance(criterion, nn.CrossEntropyLoss): for batch_idx, (inputs, targets) in enumerate(trainloader): batch_size = inputs.size(0) total += batch_size if use_cuda: inputs, targets = inputs.cuda(), targets.cuda() optimizer.zero_grad() inputs, targets = Variable(inputs), Variable(targets) outputs = net(inputs) loss = criterion(outputs, targets) loss.backward() optimizer.step() train_loss += loss.item()*batch_size _, predicted = torch.max(outputs.data, 1) correct += predicted.eq(targets.data).cpu().sum().item() elif isinstance(criterion, nn.MSELoss): for batch_idx, (inputs, targets) in enumerate(trainloader): batch_size = inputs.size(0) total += batch_size one_hot_targets = torch.FloatTensor(batch_size, 10).zero_() one_hot_targets = one_hot_targets.scatter_(1, targets.view(batch_size, 1), 1.0) one_hot_targets = one_hot_targets.float() if use_cuda: inputs, one_hot_targets = inputs.cuda(), one_hot_targets.cuda() inputs, one_hot_targets = Variable(inputs), Variable(one_hot_targets) outputs = F.softmax(net(inputs)) loss = criterion(outputs, one_hot_targets) loss.backward() optimizer.step() train_loss += loss.item()*batch_size _, predicted = torch.max(outputs.data, 1) correct += predicted.cpu().eq(targets).cpu().sum().item() return train_loss/total, 100 - 100.*correct/total def test(testloader, net, criterion, use_cuda=True): net.eval() test_loss = 0 correct = 0 total = 0 if isinstance(criterion, nn.CrossEntropyLoss): for batch_idx, (inputs, targets) in enumerate(testloader): batch_size = inputs.size(0) total += batch_size if use_cuda: inputs, targets = inputs.cuda(), targets.cuda() inputs, targets = Variable(inputs), Variable(targets) outputs = net(inputs) loss = criterion(outputs, targets) test_loss += loss.item()*batch_size _, predicted = torch.max(outputs.data, 1) correct += predicted.eq(targets.data).cpu().sum().item() elif isinstance(criterion, nn.MSELoss): for batch_idx, (inputs, targets) in enumerate(testloader): batch_size = inputs.size(0) total += batch_size one_hot_targets = torch.FloatTensor(batch_size, 10).zero_() one_hot_targets = one_hot_targets.scatter_(1, targets.view(batch_size, 1), 1.0) one_hot_targets = one_hot_targets.float() if use_cuda: inputs, one_hot_targets = inputs.cuda(), one_hot_targets.cuda() inputs, one_hot_targets = Variable(inputs), Variable(one_hot_targets) outputs = F.softmax(net(inputs)) loss = criterion(outputs, one_hot_targets) test_loss += loss.item()*batch_size _, predicted = torch.max(outputs.data, 1) correct += predicted.cpu().eq(targets).cpu().sum().item() return test_loss/total, 100 - 100.*correct/total def name_save_folder(args): save_folder = args.model + '_' + str(args.optimizer) + '_lr=' + str(args.lr) if args.lr_decay != 0.1: save_folder += '_lr_decay=' + str(args.lr_decay) save_folder += '_bs=' + str(args.batch_size) save_folder += '_wd=' + str(args.weight_decay) save_folder += '_mom=' + str(args.momentum) save_folder += '_save_epoch=' + str(args.save_epoch) if args.loss_name != 'crossentropy': save_folder += '_loss=' + str(args.loss_name) if args.noaug: save_folder += '_noaug' if args.raw_data: save_folder += '_rawdata' if args.label_corrupt_prob > 0: save_folder += '_randlabel=' + str(args.label_corrupt_prob) if args.ngpu > 1: save_folder += '_ngpu=' + str(args.ngpu) if args.idx: save_folder += '_idx=' + str(args.idx) return save_folder if __name__ == '__main__': # Training options parser = argparse.ArgumentParser(description='PyTorch CIFAR10 Training') parser.add_argument('--batch_size', default=128, type=int) parser.add_argument('--lr', default=0.1, type=float, help='learning rate') parser.add_argument('--lr_decay', default=0.1, type=float, help='learning rate decay rate') parser.add_argument('--optimizer', default='sgd', help='optimizer: sgd | adam') parser.add_argument('--weight_decay', default=0.0005, type=float) parser.add_argument('--momentum', default=0.9, type=float) parser.add_argument('--epochs', default=300, type=int, metavar='N', help='number of total epochs to run') parser.add_argument('--save', default='trained_nets',help='path to save trained nets') p [truncated — 5595 more characters] ``` ### loss-landscape/model_loader.py ```python import os import cifar10.model_loader def load(dataset, model_name, model_file, data_parallel=False): if dataset == 'cifar10': net = cifar10.model_loader.load(model_name, model_file, data_parallel) return net ``` ### loss-landscape/h5_util.py ```python """ Serialization and deserialization of directions in the direction file. """ import torch def write_list(f, name, direction): """ Save the direction to the hdf5 file with name as the key Args: f: h5py file object name: key name_surface_file direction: a list of tensors """ grp = f.create_group(name) for i, l in enumerate(direction): if isinstance(l, torch.Tensor): l = l.numpy() grp.create_dataset(str(i), data=l) def read_list(f, name): """ Read group with name as the key from the hdf5 file and return a list numpy vectors. """ grp = f[name] return [grp[str(i)] for i in range(len(grp))] ``` ### loss-landscape/evaluation.py ```python """ The calculation to be performed at each point (modified model), evaluating the loss value, accuracy and eigen values of the hessian matrix """ import torch import torch.nn as nn import torch.nn.functional as F import time from torch.autograd.variable import Variable def eval_loss(net, criterion, loader, use_cuda=False): """ Evaluate the loss value for a given 'net' on the dataset provided by the loader. Args: net: the neural net model criterion: loss function loader: dataloader use_cuda: use cuda or not Returns: loss value and accuracy """ correct = 0 total_loss = 0 total = 0 # number of samples num_batch = len(loader) if use_cuda: net.cuda() net.eval() with torch.no_grad(): if isinstance(criterion, nn.CrossEntropyLoss): for batch_idx, (inputs, targets) in enumerate(loader): batch_size = inputs.size(0) total += batch_size inputs = Variable(inputs) targets = Variable(targets) if use_cuda: inputs, targets = inputs.cuda(), targets.cuda() outputs = net(inputs) loss = criterion(outputs, targets) total_loss += loss.item()*batch_size _, predicted = torch.max(outputs.data, 1) correct += predicted.eq(targets).sum().item() elif isinstance(criterion, nn.MSELoss): for batch_idx, (inputs, targets) in enumerate(loader): batch_size = inputs.size(0) total += batch_size inputs = Variable(inputs) one_hot_targets = torch.FloatTensor(batch_size, 10).zero_() one_hot_targets = one_hot_targets.scatter_(1, targets.view(batch_size, 1), 1.0) one_hot_targets = one_hot_targets.float() one_hot_targets = Variable(one_hot_targets) if use_cuda: inputs, one_hot_targets = inputs.cuda(), one_hot_targets.cuda() outputs = F.softmax(net(inputs)) loss = criterion(outputs, one_hot_targets) total_loss += loss.item()*batch_size _, predicted = torch.max(outputs.data, 1) correct += predicted.cpu().eq(targets).sum().item() return total_loss/total, 100.*correct/total ``` ### loss-landscape/plot_trajectory.py ```python """ Plot the optimization path in the space spanned by principle directions. """ import numpy as np import torch import copy import math import h5py import os import argparse import model_loader import net_plotter from projection import setup_PCA_directions, project_trajectory import plot_2D if __name__ == '__main__': parser = argparse.ArgumentParser(description='Plot optimization trajectory') parser.add_argument('--dataset', default='cifar10', help='dataset') parser.add_argument('--model', default='resnet56', help='trained models') parser.add_argument('--model_folder', default='', help='folders for models to be projected') parser.add_argument('--dir_type', default='weights', help="""direction type: weights (all weights except bias and BN paras) | states (include BN.running_mean/var)""") parser.add_argument('--ignore', default='', help='ignore bias and BN paras: biasbn (no bias or bn)') parser.add_argument('--prefix', default='model_', help='prefix for the checkpint model') parser.add_argument('--suffix', default='.t7', help='prefix for the checkpint model') parser.add_argument('--start_epoch', default=0, type=int, help='min index of epochs') parser.add_argument('--max_epoch', default=300, type=int, help='max number of epochs') parser.add_argument('--save_epoch', default=1, type=int, help='save models every few epochs') parser.add_argument('--dir_file', default='', help='load the direction file for projection') args = parser.parse_args() #-------------------------------------------------------------------------- # load the final model #-------------------------------------------------------------------------- last_model_file = args.model_folder + '/' + args.prefix + str(args.max_epoch) + args.suffix net = model_loader.load(args.dataset, args.model, last_model_file) w = net_plotter.get_weights(net) s = net.state_dict() #-------------------------------------------------------------------------- # collect models to be projected #-------------------------------------------------------------------------- model_files = [] for epoch in range(args.start_epoch, args.max_epoch + args.save_epoch, args.save_epoch): model_file = args.model_folder + '/' + args.prefix + str(epoch) + args.suffix assert os.path.exists(model_file), 'model %s does not exist' % model_file model_files.append(model_file) #-------------------------------------------------------------------------- # load or create projection directions #-------------------------------------------------------------------------- if args.dir_file: dir_file = args.dir_file else: dir_file = setup_PCA_directions(args, model_files, w, s) #-------------------------------------------------------------------------- # projection trajectory to given directions #-------------------------------------------------------------------------- proj_file = project_trajectory(dir_file, w, s, args.dataset, args.model, model_files, args.dir_type, 'cos') plot_2D.plot_trajectory(proj_file, dir_file) ``` ### loss-landscape/scheduler.py ```python """ A task scheduler that assign unfinished jobs to different workers. """ import numpy as np def get_unplotted_indices(vals, xcoordinates, ycoordinates=None): """ Args: vals: values at (x, y), with value -1 when the value is not yet calculated. xcoordinates: x locations, i.e.,[-1, -0.5, 0, 0.5, 1] ycoordinates: y locations, i.e.,[-1, -0.5, 0, 0.5, 1] Returns: - a list of indices into vals for points that have not yet been calculated. - a list of corresponding coordinates, with one x/y coordinate per row. """ # Create a list of indices into the vectorizes vals inds = np.array(range(vals.size)) # Select the indices of the un-recorded entries, assuming un-recorded entries # will be smaller than zero. In case some vals (other than loss values) are # negative and those indexces will be selected again and calcualted over and over. inds = inds[vals.ravel() <= 0] # Make lists containing the x- and y-coodinates of the points to be plotted if ycoordinates is not None: # If the plot is 2D, then use meshgrid to enumerate all coordinates in the 2D mesh xcoord_mesh, ycoord_mesh = np.meshgrid(xcoordinates, ycoordinates) s1 = xcoord_mesh.ravel()[inds] s2 = ycoord_mesh.ravel()[inds] return inds, np.c_[s1,s2] else: return inds, xcoordinates.ravel()[inds] def split_inds(num_inds, nproc): """ Evenly slice out a set of jobs that are handled by each MPI process. - Assuming each job takes the same amount of time. - Each process handles an (approx) equal size slice of jobs. - If the number of processes is larger than rows to divide up, then some high-rank processes will receive an empty slice rows, e.g., there will be 3, 2, 2, 2 jobs assigned to rank0, rank1, rank2, rank3 given 9 jobs with 4 MPI processes. """ chunk = num_inds // nproc remainder = num_inds % nproc splitted_idx = [] for rank in range(0, nproc): # Set the starting index for this slice start_idx = rank * chunk + min(rank, remainder) # The stopping index can't go beyond the end of the array stop_idx = start_idx + chunk + (rank < remainder) splitted_idx.append(range(start_idx, stop_idx)) return splitted_idx def get_job_indices(vals, xcoordinates, ycoordinates, comm): """ Prepare the job indices over which coordinate to calculate. Args: vals: the value matrix xcoordinates: x locations, i.e.,[-1, -0.5, 0, 0.5, 1] ycoordinates: y locations, i.e.,[-1, -0.5, 0, 0.5, 1] comm: MPI environment Returns: inds: indices that splitted for current rank coords: coordinates for current rank inds_nums: max number of indices for all ranks """ inds, coords = get_unplotted_indices(vals, xcoordinates, ycoordinates) rank = 0 if comm is None else comm.Get_rank() nproc = 1 if comm is None else comm.Get_size() splitted_idx = split_inds(len(inds), nproc) # Split the indices over the available MPI processes inds = inds[splitted_idx[rank]] coords = coords[splitted_idx[rank]] # Figure out the number of jobs that each MPI process needs to calculate. inds_nums = [len(idx) for idx in splitted_idx] return inds, coords, inds_nums ``` ### loss-landscape/mpi4pytorch.py ```python """ mpi4pytorch.py This module contains convenience methods that make it easy to use mpi4py. The available functions handle memory allocation and other data formatting tasks so that tensors can be easily reduced/broadcast using 1 line of code. """ import numpy as np import mpi4py def setup_MPI(): try: from mpi4py import MPI comm = MPI.COMM_WORLD # Convert the Object to a Class so that it is possible to add attributes later class A(mpi4py.MPI.Intracomm): pass comm = A(comm) except: comm = None return comm def print_once(comm, *message): if not comm or comm.Get_rank()==0: print (''.join(str(i) for i in message)) def is_master(comm): return not comm or comm.Get_rank()==0 def allreduce_max(comm, array, display_info=False): if not comm: return array array = np.asarray(array, dtype='d') total = np.zeros_like(array) float_min = np.finfo(np.float).min total.fill(float_min) if display_info: print ("(%d): sum=%f : size=%d"%(get_rank(comm), np.sum(array), array.nbytes)) rows = str(comm.gather(array.shape[0])) cols = str(comm.gather(array.shape[1])) print_once(comm, "reduce: %s, %s"%(rows, cols)) comm.Allreduce(array, total, op=mpi4py.MPI.MAX) return total def allreduce_min(comm, array, display_info=False): if not comm: return array array = np.asarray(array, dtype='d') total = np.zeros_like(array) float_max = np.finfo(np.float).max total.fill(float_max) if display_info: print ("(%d): sum=%f : size=%d"%(get_rank(comm), np.sum(array), array.nbytes)) rows = str(comm.gather(array.shape[0])) cols = str(comm.gather(array.shape[1])) print_once(comm, "reduce: %s, %s"%(rows, cols)) comm.Allreduce(array, total, op=mpi4py.MPI.MIN) return total def reduce_max(comm, array, display_info=False): if not comm: return array array = np.asarray(array, dtype='d') total = np.zeros_like(array) float_min = np.finfo(np.float).min total.fill(float_min) if display_info: print ("(%d): sum=%f : size=%d"%(get_rank(comm), np.sum(array), array.nbytes)) rows = str(comm.gather(array.shape[0])) cols = str(comm.gather(array.shape[1])) print_once(comm, "reduce: %s, %s"%(rows, cols)) comm.Reduce(array, total, op=mpi4py.MPI.MAX, root=0) return total def reduce_min(comm, array, display_info=False): if not comm: return array array = np.asarray(array, dtype='d') total = np.zeros_like(array) float_max = np.finfo(np.float).max total.fill(float_max) if display_info: print ("(%d): sum=%f : size=%d"%(get_rank(comm), np.sum(array), array.nbytes)) rows = str(comm.gather(array.shape[0])) cols = str(comm.gather(array.shape[1])) print_once(comm, "reduce: %s, %s"%(rows, cols)) comm.Reduce(array, total, op=mpi4py.MPI.MIN, root=0) return total def barrier(comm): if not comm: return comm.barrier() def get_mpi_info(): try: return mpi4py.MPI.get_vendor() except ImportError: return "none" def get_rank(comm): try: return comm.Get_rank() except ImportError: return 0 def get_num_procs(comm): try: return comm.Get_size() except ImportError: return 1 ``` ### loss-landscape/dataloader.py ```python import torch import torchvision from torchvision import transforms import os import numpy as np import argparse def get_relative_path(file): script_dir = os.path.dirname(__file__) # <-- absolute dir the script is in return os.path.join(script_dir, file) def load_dataset(dataset='cifar10', datapath='cifar10/data', batch_size=128, \ threads=2, raw_data=False, data_split=1, split_idx=0, \ trainloader_path="", testloader_path=""): """ Setup dataloader. The data is not randomly cropped as in training because of we want to esimate the loss value with a fixed dataset. Args: raw_data: raw images, no data preprocessing data_split: the number of splits for the training dataloader split_idx: the index for the split of the dataloader, starting at 0 Returns: train_loader, test_loader """ # use specific dataloaders if trainloader_path and testloader_path: assert os.path.exists(trainloader_path), 'trainloader does not exist' assert os.path.exists(testloader_path), 'testloader does not exist' train_loader = torch.load(trainloader_path) test_loader = torch.load(testloader_path) return train_loader, test_loader assert split_idx < data_split, 'the index of data partition should be smaller than the total number of split' if dataset == 'cifar10': normalize = transforms.Normalize(mean=[x/255.0 for x in [125.3, 123.0, 113.9]], std=[x/255.0 for x in [63.0, 62.1, 66.7]]) data_folder = get_relative_path(datapath) if raw_data: transform = transforms.Compose([ transforms.ToTensor() ]) else: transform = transforms.Compose([ transforms.ToTensor(), normalize, ]) trainset = torchvision.datasets.CIFAR10(root=data_folder, train=True, download=True, transform=transform) # If data_split>1, then randomly select a subset of the data. E.g., if datasplit=3, then # randomly choose 1/3 of the data. if data_split > 1: indices = torch.tensor(np.arange(len(trainset))) data_num = len(trainset) // data_split # the number of data in a chunk of the split # Randomly sample indices. Use seed=0 in the generator to make this reproducible state = np.random.get_state() np.random.seed(0) indices = np.random.choice(indices, data_num, replace=False) np.random.set_state(state) train_sampler = torch.utils.data.sampler.SubsetRandomSampler(indices) train_loader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, sampler=train_sampler, shuffle=False, num_workers=threads) else: kwargs = {'num_workers': 2, 'pin_memory': True} train_loader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, shuffle=False, **kwargs) testset = torchvision.datasets.CIFAR10(root=data_folder, train=False, download=False, transform=transform) test_loader = torch.utils.data.DataLoader(testset, batch_size=batch_size, shuffle=False, num_workers=threads) return train_loader, test_loader ############################################################### #### MAIN ############################################################### if __name__ == '__main__': parser = argparse.ArgumentParser(description='PyTorch CIFAR10 Training') parser.add_argument('--mpi', '-m', action='store_true', help='use mpi') parser.add_argument('--cuda', '-c', action='store_true', help='use cuda') parser.add_argument('--threads', default=2, type=int, help='number of threads') parser.add_argument('--batch_size', default=128, type=int, help='minibatch size') parser.add_argument('--dataset', default='cifar10', help='cifar10 | imagenet') parser.add_argument('--datapath', default='cifar10/data', metavar='DIR', help='path to the dataset') parser.add_argument('--raw_data', action='store_true', default=False, help='do not normalize data') parser.add_argument('--data_split', default=1, type=int, help='the number of splits for the dataloader') parser.add_argument('--split_idx', default=0, type=int, help='the index of data splits for the dataloader') parser.add_argument('--trainloader', default='', help='path to the dataloader with random labels') parser.add_argument('--testloader', default='', help='path to the testloader with random labels') args = parser.parse_args() trainloader, testloader = load_dataset(args.dataset, args.datapath, args.batch_size, args.threads, args.raw_data, args.data_split, args.split_idx, args.trainloader, args.testloader) print('num of batches: %d' % len(trainloader)) for batch_idx, (inputs, targets) in enumerate(trainloader): print('batch_idx: %d batch_size: %d'%(batch_idx, len(inputs))) ``` ### loss-landscape/hess_vec_prod.py ```python import torch import time import numpy as np from torch import nn from torch.autograd import Variable from scipy.sparse.linalg import LinearOperator, eigsh ################################################################################ # Supporting Functions ################################################################################ def npvec_to_tensorlist(vec, params): """ Convert a numpy vector to a list of tensor with the same dimensions as params Args: vec: a 1D numpy vector params: a list of parameters from net Returns: rval: a list of tensors with the same shape as params """ loc = 0 rval = [] for p in params: numel = p.data.numel() rval.append(torch.from_numpy(vec[loc:loc+numel]).view(p.data.shape).float()) loc += numel assert loc == vec.size, 'The vector has more elements than the net has parameters' return rval def gradtensor_to_npvec(net, include_bn=False): """ Extract gradients from net, and return a concatenated numpy vector. Args: net: trained model include_bn: If include_bn, then gradients w.r.t. BN parameters and bias values are also included. Otherwise only gradients with dim > 1 are considered. Returns: a concatenated numpy vector containing all gradients """ filter = lambda p: include_bn or len(p.data.size()) > 1 return np.concatenate([p.grad.data.cpu().numpy().ravel() for p in net.parameters() if filter(p)]) ################################################################################ # For computing Hessian-vector products ################################################################################ def eval_hess_vec_prod(vec, params, net, criterion, dataloader, use_cuda=False): """ Evaluate product of the Hessian of the loss function with a direction vector "vec". The product result is saved in the grad of net. Args: vec: a list of tensor with the same dimensions as "params". params: the parameter list of the net (ignoring biases and BN parameters). net: model with trained parameters. criterion: loss function. dataloader: dataloader for the dataset. use_cuda: use GPU. """ if use_cuda: net.cuda() vec = [v.cuda() for v in vec] net.eval() net.zero_grad() # clears grad for every parameter in the net for batch_idx, (inputs, targets) in enumerate(dataloader): inputs, targets = Variable(inputs), Variable(targets) if use_cuda: inputs, targets = inputs.cuda(), targets.cuda() outputs = net(inputs) loss = criterion(outputs, targets) grad_f = torch.autograd.grad(loss, inputs=params, create_graph=True) # Compute inner product of gradient with the direction vector prod = Variable(torch.zeros(1)).type(type(grad_f[0].data)) for (g, v) in zip(grad_f, vec): prod = prod + (g * v).cpu().sum() # Compute the Hessian-vector product, H*v # prod.backward() computes dprod/dparams for every parameter in params and # accumulate the gradients into the params.grad attributes prod.backward() ################################################################################ # For computing Eigenvalues of Hessian ################################################################################ def min_max_hessian_eigs(net, dataloader, criterion, rank=0, use_cuda=False, verbose=False): """ Compute the largest and the smallest eigenvalues of the Hessian marix. Args: net: the trained model. dataloader: dataloader for the dataset, may use a subset of it. criterion: loss function. rank: rank of the working node. use_cuda: use GPU verbose: print more information Returns: maxeig: max eigenvalue mineig: min eigenvalue hess_vec_prod.count: number of iterations for calculating max and min eigenvalues """ params = [p for p in net.parameters() if len(p.size()) > 1] N = sum(p.numel() for p in params) def hess_vec_prod(vec): hess_vec_prod.count += 1 # simulates a static variable vec = npvec_to_tensorlist(vec, params) start_time = time.time() eval_hess_vec_prod(vec, params, net, criterion, dataloader, use_cuda) prod_time = time.time() - start_time if verbose and rank == 0: print(" Iter: %d time: %f" % (hess_vec_prod.count, prod_time)) return gradtensor_to_npvec(net) hess_vec_prod.count = 0 if verbose and rank == 0: print("Rank %d: computing max eigenvalue" % rank) A = LinearOperator((N, N), matvec=hess_vec_prod) eigvals, eigvecs = eigsh(A, k=1, tol=1e-2) maxeig = eigvals[0] if verbose and rank == 0: print('max eigenvalue = %f' % maxeig) # If the largest eigenvalue is positive, shift matrix so that any negative eigenvalue is now the largest # We assume the smallest eigenvalue is zero or less, and so this shift is more than what we need shift = maxeig*.51 def shifted_hess_vec_prod(vec): return hess_vec_prod(vec) - shift*vec if verbose and rank == 0: print("Rank %d: Computing shifted eigenvalue" % rank) A = LinearOperator((N, N), matvec=shifted_hess_vec_prod) eigvals, eigvecs = eigsh(A, k=1, tol=1e-2) eigvals = eigvals + shift mineig = eigvals[0] if verbose and rank == 0: print('min eigenvalue = ' + str(mineig)) if maxeig <= 0 and mineig > 0: maxeig, mineig = mineig, maxeig return maxeig, mineig, hess_vec_prod.count ``` [28 more indexed source files omitted to keep this export small. The full file list is in the Codebase structure section above.]