#--------------------------------------------------------------------------
# Setup dataloader
#--------------------------------------------------------------------------
# download CIFAR10 if it does not exit
if rank == 0 and args.dataset == 'cifar10':
torchvision.datasets.CIFAR10(root=args.dataset + '/data', train=True, download=True)
mpi4pytorch.barrier(comm)
trainloader, testloader = dataloader.load_dataset(args.dataset, args.datapath,
args.batch_size, args.threads, args.raw_data,
args.data_split, args.split_idx,
args.trainloader, args.testloader)
#--------------------------------------------------------------------------
# Start the computation
#--------------------------------------------------------------------------
crunch(surf_file, net, w, s, d, trainloader, 'train_loss', 'train_acc', comm, rank, args)
# crunch(surf_file, net, w, s, d, testloader, 'test_loss', 'test_acc', comm, rank, args)
#--------------------------------------------------------------------------
# Plot figures
#--------------------------------------------------------------------------
if args.plot and rank == 0:
if args.y and args.proj_file:
plot_2D.plot_contour_trajectory(surf_file, dir_file, args.proj_file, 'train_loss', args.show)
elif args.y:
plot_2D.plot_2d_contour(surf_file, 'train_loss', args.vmin, args.vmax, args.vlevel, args.show)
else:
plot_1D.plot_1d_loss_err(surf_file, args.xmin, args.xmax, args.loss_max, args.log, args.show)
download=True)
mpi.barrier(comm)
trainloader, testloader = dataloader.load_dataset(
args.dataset, args.datapath, args.batch_size, args.threads,
args.raw_data, args.data_split, args.split_idx, args.trainloader,
args.testloader)
#--------------------------------------------------------------------------
# Start the computation
#--------------------------------------------------------------------------
crunch(surf_file, net, w, s, d, trainloader, 'train_loss', 'train_acc',
comm, rank, args)
# crunch(surf_file, net, w, s, d, testloader, 'test_loss', 'test_acc', comm, rank, args)
#--------------------------------------------------------------------------
# Plot figures
#--------------------------------------------------------------------------
if args.plot and rank == 0:
if args.y and args.proj_file:
plot_2D.plot_contour_trajectory(surf_file, dir_file,
args.proj_file, 'train_loss',
args.show)
elif args.y:
plot_2D.plot_2d_contour(surf_file, 'train_loss', args.vmin,
args.vmax, args.vlevel, args.show)
else:
plot_1D.plot_1d_loss_err(surf_file, args.xmin, args.xmax,
args.loss_max, args.log, args.show)
'================================================================================='
)
# --------------------------------------------------------------------------
# Plot figures
# --------------------------------------------------------------------------
if args.plot and rank == 0:
start_time = time.time()
if args.y and args.proj_file: # if args contains y and args contains proj_file. Often for 2d plots, 3 models.
plot_2D.plot_contour_trajectory(surf_file, dir_file,
args.proj_file, f'{dataset}_loss',
args.show)
print(
'----------------------------\n ----plot 2d countour and trajectory ------\n----------------------'
)
elif args.y: # if args contains y, while NOT contains proj_file. Often for 2d plot, 1 model
plot_2D.plot_2d_contour(surf_file, f'{dataset}_loss', args.vmin,
args.vmax, args.vlevel, args.show)
print(
'----------------------------\n ----plot 2d countour ------\n----------------------'
)
else: # if args does NOT contains y; must be 1d plot, 2 models.
plot_1D.plot_1d_loss_err(surf_file, args.xmin, args.xmax,
args.loss_max, args.log, args.show)
print(
'----------------------------\n ----plot 1d interpolation ------\n----------------------'
)
plot_time = time.time() - start_time
print('Plotting time:', plot_time)
# save all args in a txt file in the same folder
print_all_args = 1
