Request for Landscape Implementation Code

[lzgithub2002]

Hi, thank you for your excellent work! Could you please consider providing an implementation or example code for the landscape visualization？

Thanks!

[effl-lab]

Hi---sorry for the very late reply. the first author has graduated and he might have missed this issue. @sangyoon-lee99 can you handle this?

[sangyoon-lee99]

Hi @lzgithub2002 , thank you for raising this issue.
I apologize for the delay in our response.

Our implementation is based on the original loss-landscape source code (https://github.com/tomgoldstein/loss-landscape).
We've introduced a few modifications to adapt the landscape visualization for an image fitting task, along with computing the direction derived from the Hessian.

Here is a summary of the key changes:

net_plotter.setup_direction ()

### MODIFIED : Add a 'loader' argument for Hessian-based direction computation.
def setup_direction(args, dir_file, net, loader):
    ...
    # Create the plotting directions
    f = h5py.File(dir_file,'w') # create file, fail if exists
    if not args.dir_file:
        print("Setting up the plotting directions...")

        ### MODIFIED : Set x-direction based on the Hessian's top eigenvector.
        if args.hessdir:
            inputs, targets = next(iter(loader))
            if args.model_file2:
                net2 = model_loader.load(args.dataset, args.model, args.model_file2)
                ydirection = create_target_direction(net, net2, args.dir_type)
            net = net.cuda()
            
            ### Use PyHessian library to compute the top eigenvector
            hessian_comp = hessian(net, nn.MSELoss(), data=(inputs, targets), cuda=True)
            
            top_eigenvalues, top_eigenvector = hessian_comp.eigenvalues(top_n=1)
            xdirection = top_eigenvector[0]
            
            h5_util.write_list(f, 'xdirection', xdirection)
            h5_util.write_list(f, 'ydirection', ydirection)
            
        ### The rest of the function remains the same as the original.
        else:   
            if args.model_file2:
            ...

plot_surface.py (main script)

...
### MODIFIED: Use only a train loader for our image fitting dataset. (same as our main code!)
loaded_data = ImageFitting(512, args.dataset)
trainloader = DataLoader(loaded_data, batch_size=2**18, shuffle=False, pin_memory=True, num_workers=16)

... 

### MODIFIED: Removed accuracy-related arguments from the "crunch" function call.
crunch(surf_file, net, w, s, d, trainloader, 'train_loss', comm, rank, args)
...

We believe these changes address the question about our implementation. Please let us know if you have any further questions or encounter any problems.

Thank you!