shuffle_train=False, ) model = model_cfg.base(*model_cfg.args, num_classes=num_classes, **model_cfg.kwargs) model.cuda() print("Loading model %s" % args.file) checkpoint = torch.load(args.file) model.load_state_dict(checkpoint["state_dict"]) criterion = torch.nn.CrossEntropyLoss() if args.use_test: loader = loaders["test"] else: loader = loaders["train"] max_eval, min_eval, hvps, pos_evals, neg_evals = min_max_hessian_eigs( model, loader, criterion, use_cuda=True, verbose=True) print("Maximum eigenvalue: ", max_eval) print("Minimum eigenvalue: ", min_eval) print("Number of Hessian vector products: ", hvps) print("Saving all eigenvalues to ", args.save_path) np.savez(args.save_path, pos_evals=pos_evals.cpu().numpy(), neg_evals=neg_evals.cpu().numpy())
def crunch_hessian_eigs(surf_file, net, w, s, d, dataloader, comm, rank, args): """ Calculate eigen values of the hessian matrix of a given model in parallel using mpi reduce. This is the synchronized version. """ f = h5py.File(surf_file, 'r+' if rank == 0 else 'r') min_eig, max_eig = [], [] xcoordinates = f['xcoordinates'][:] ycoordinates = f['ycoordinates'][:] if 'ycoordinates' in f.keys() else None if 'min_eig' not in f.keys(): shape = xcoordinates.shape if ycoordinates is None else (len(xcoordinates), len(ycoordinates)) max_eig = -np.ones(shape=shape) min_eig = np.ones(shape=shape) if rank == 0: f['min_eig'] = min_eig f['max_eig'] = max_eig else: min_eig = f['min_eig'][:] max_eig = f['max_eig'][:] # Generate a list of all indices that need to be filled in. # The coordinates of each unfilled index are stored in 'coords'. inds, coords, inds_nums = scheduler.get_job_indices(max_eig, xcoordinates, ycoordinates, comm) print('Computing %d values for rank %d' % (len(inds), rank)) criterion = nn.CrossEntropyLoss() # set the loss function criteria # Loop over all un-calculated coords start_time = time.time() total_sync = 0.0 for count, ind in enumerate(inds): # Get the coordinates of the points being calculated coord = coords[count] # Load the weights corresponding to those coordinates into the net if args.dir_type == 'weights': net_plotter.set_weights(net.module if args.ngpu > 1 else net, w, d, coord) elif args.dir_type == 'states': net_plotter.set_states(net.module if args.ngpu > 1 else net, s, d, coord) # Compute the eign values of the hessian matrix compute_start = time.time() maxeig, mineig, iter_count = hess_vec_prod.min_max_hessian_eigs(net, dataloader, \ criterion, rank=rank, use_cuda=args.cuda, verbose=True) compute_time = time.time() - compute_start # Record the result in the local array max_eig.ravel()[ind] = maxeig min_eig.ravel()[ind] = mineig # Send updated plot data to the master node sync_start_time = time.time() max_eig = mpi4pytorch.reduce_max(comm, max_eig) min_eig = mpi4pytorch.reduce_min(comm, min_eig) sync_time = time.time() - sync_start_time total_sync += sync_time # Only the master node writes to the file - this avoids write conflicts if rank == 0: f['max_eig'][:] = max_eig f['min_eig'][:] = min_eig print("rank: %d %d/%d (%0.2f%%) %d\t %s \tmaxeig:%8.5f \tmineig:%8.5f \titer: %d \ttime:%.2f \tsync:%.2f" % ( \ rank, count + 1, len(inds), 100.0 * (count + 1) / len(inds), ind, str(coord), \ maxeig, mineig, iter_count, compute_time, sync_time)) # This is only needed to make MPI run smoothly. If this process has less work # than the rank0 process, then we need to keep calling allreduce so the rank0 process doesn't block for i in range(max(inds_nums) - len(inds)): max_eig = mpi4pytorch.reduce_max(comm, max_eig) min_eig = mpi4pytorch.reduce_min(comm, min_eig) total_time = time.time() - start_time print('Rank %d done! Total time: %f Sync: %f ' % (rank, total_time, total_sync)) f.close()
def crunch_hessian_eigs(surf_file, net, w, s, d, dataloader, comm, rank, args): """ Calculate eigen values of the hessian matrix of a given model in parallel using mpi reduce. This is the synchronized version. """ f = h5py.File(surf_file, 'r+' if rank == 0 else 'r') min_eig, max_eig = [], [] xcoordinates = f['xcoordinates'][:] ycoordinates = f['ycoordinates'][:] if 'ycoordinates' in f.keys() else None if 'min_eig' not in f.keys(): shape = xcoordinates.shape if ycoordinates is None else (len(xcoordinates),len(ycoordinates)) max_eig = -np.ones(shape=shape) min_eig = np.ones(shape=shape) if rank == 0: f['min_eig'] = min_eig f['max_eig'] = max_eig else: min_eig = f['min_eig'][:] max_eig = f['max_eig'][:] # Generate a list of all indices that need to be filled in. # The coordinates of each unfilled index are stored in 'coords'. inds, coords, inds_nums = scheduler.get_job_indices(max_eig, xcoordinates, ycoordinates, comm) print('Computing %d values for rank %d'% (len(inds), rank)) criterion = nn.CrossEntropyLoss() # set the loss function criteria # Loop over all un-calculated coords start_time = time.time() total_sync = 0.0 for count, ind in enumerate(inds): # Get the coordinates of the points being calculated coord = coords[count] # Load the weights corresponding to those coordinates into the net if args.dir_type == 'weights': net_plotter.set_weights(net.module if args.ngpu > 1 else net, w, d, coord) elif args.dir_type == 'states': net_plotter.set_states(net.module if args.ngpu > 1 else net, s, d, coord) # Compute the eign values of the hessian matrix compute_start = time.time() maxeig, mineig, iter_count = hess_vec_prod.min_max_hessian_eigs(net, dataloader, \ criterion, rank=rank, use_cuda=args.cuda, verbose=True) compute_time = time.time() - compute_start # Record the result in the local array max_eig.ravel()[ind] = maxeig min_eig.ravel()[ind] = mineig # Send updated plot data to the master node sync_start_time = time.time() max_eig = mpi4pytorch.reduce_max(comm, max_eig) min_eig = mpi4pytorch.reduce_min(comm, min_eig) sync_time = time.time() - sync_start_time total_sync += sync_time # Only the master node writes to the file - this avoids write conflicts if rank == 0: f['max_eig'][:] = max_eig f['min_eig'][:] = min_eig print("rank: %d %d/%d (%0.2f%%) %d\t %s \tmaxeig:%8.5f \tmineig:%8.5f \titer: %d \ttime:%.2f \tsync:%.2f" % ( \ rank, count + 1, len(inds), 100.0 * (count + 1)/len(inds), ind, str(coord), \ maxeig, mineig, iter_count, compute_time, sync_time)) # This is only needed to make MPI run smoothly. If this process has less work # than the rank0 process, then we need to keep calling allreduce so the rank0 process doesn't block for i in range(max(inds_nums) - len(inds)): max_eig = mpi4pytorch.reduce_max(comm, max_eig) min_eig = mpi4pytorch.reduce_min(comm, min_eig) total_time = time.time() - start_time print('Rank %d done! Total time: %f Sync: %f '%(rank, total_time, total_sync)) f.close()