# turn model to eval mode
model.eval()
if batch_num == 1:
hessian_comp = hessian(model,
criterion,
data=hessian_dataloader,
cuda=args.cuda,
record_data=True)
else:
hessian_comp = hessian(model,
criterion,
dataloader=hessian_dataloader,
cuda=args.cuda,
record_data=True)
print(
'********** finish data loading and begin Hessian computation **********')
print("Computing eigenvalues...")
top_eigenvalues, _ = hessian_comp.eigenvalues(top_n=20, debug=True)
print("Computing trace...")
trace = hessian_comp.trace(debug=True)
print("Computing ESD...")
density_eigen, density_weight = hessian_comp.density(debug=True)
print('\n***Top Eigenvalues: ', top_eigenvalues)
print('\n***Trace: ', np.mean(trace))
get_esd_plot(density_eigen, density_weight)
f_density_eigen.write(str(density_eigen[0][i]) + '\n')
f_density_eigen.close()
if not os.path.exists(args.output + "/density_weight"):
os.makedirs(args.output + "/density_weight")
f_density_weight = open(
args.output + '/density_weight' + '/density_weight' + index + '.txt',
'w')
for i in range(len(density_weight[0])):
f_density_weight.write(str(density_weight[0][i]) + '\n')
f_density_weight.close()
if not os.path.exists(args.output + "/esd_plot"):
os.makedirs(args.output + "/esd_plot")
get_esd_plot(density_eigen,
density_weight,
path=args.output + '/esd_plot' + '/esd_plot' + index + '.pdf')
print("\n***finish " + index + "\n\n")
if args.compute_each_layer:
layer = -1
for param in model.parameters():
layer += 1
if not param.requires_grad:
continue
if batch_num == 1:
hessian_comp = hessian(model,
criterion,
data=hessian_dataloader,
cuda=args.cuda,
whole_model=False,