some reformatting; progress logging; dataparallel check

Author: noahgolmant

example/main.py

@@ -14,41 +14,69 @@
 
 
 def extra_args(parser):
- parser.add_argument('--num_eigenthings', default=5, type=int,
- help='number of eigenvals/vecs to compute')
- parser.add_argument('--batch_size', default=128, type=int,
- help='train set batch size')
- parser.add_argument('--eval_batch_size', default=16, type=int,
- help='test set batch size')
- parser.add_argument('--momentum', default=0.0, type=float,
- help='power iteration momentum term')
- parser.add_argument('--num_steps', default=20, type=int,
- help='number of power iter steps')
- parser.add_argument('--cuda', action='store_true',
- help='if true, use CUDA/GPUs')
+ parser.add_argument(
+ "--num_eigenthings",
+ default=5,
+ type=int,
+ help="number of eigenvals/vecs to compute",
+ )
+ parser.add_argument(
+ "--batch_size", default=128, type=int, help="train set batch size"
+ )
+ parser.add_argument(
+ "--eval_batch_size", default=16, type=int, help="test set batch size"
+ )
+ parser.add_argument(
+ "--momentum", default=0.0, type=float, help="power iteration momentum term"
+ )
+ parser.add_argument(
+ "--num_steps", default=50, type=int, help="number of power iter steps"
+ )
+ parser.add_argument("--max_samples", default=2048, type=int)
+ parser.add_argument("--cuda", action="store_true", help="if true, use CUDA/GPUs")
+ parser.add_argument(
+ "--full_dataset",
+ action="store_true",
+ help="if true,\
+ loop over all batches in set for each gradient step",
+ )
+ parser.add_argument("--fname", default="", type=str)
+ parser.add_argument("--mode", type=str, choices=["power_iter", "lanczos"])
 
 
 def main(args):
- trainloader, testloader = build_dataset('cifar10',
- dataroot=args.dataroot,
- batch_size=args.batch_size,
- eval_batch_size=args.eval_batch_size,
- num_workers=2)
- model = build_model('ResNet18', num_classes=10)
+ trainloader, testloader = build_dataset(
+ "cifar10",
+ dataroot=args.dataroot,
+ batch_size=args.batch_size,
+ eval_batch_size=args.eval_batch_size,
+ num_workers=2,
+ )
+ if args.fname:
+ print("Loading model from %s" % args.fname)
+ model = torch.load(args.fname, map_location="cpu").cuda()
+ else:
+ model = build_model("ResNet18", num_classes=10)
  criterion = torch.nn.CrossEntropyLoss()
- eigenvals, eigenvecs = compute_hessian_eigenthings(model, testloader,
- criterion,
- args.num_eigenthings,
- args.num_steps,
- momentum=args.momentum,
- use_gpu=args.cuda)
+ eigenvals, eigenvecs = compute_hessian_eigenthings(
+ model,
+ testloader,
+ criterion,
+ args.num_eigenthings,
+ mode=args.mode,
+ # power_iter_steps=args.num_steps,
+ max_samples=args.max_samples,
+ # momentum=args.momentum,
+ full_dataset=args.full_dataset,
+ use_gpu=args.cuda,
+ )
  print("Eigenvecs:")
  print(eigenvecs)
  print("Eigenvals:")
  print(eigenvals)
- track.metric(iteration=0, eigenvals=eigenvals)
+ # track.metric(iteration=0, eigenvals=eigenvals)
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
  skeletor.supply_args(extra_args)
  skeletor.execute(main)

hessian_eigenthings/__init__.py

@@ -1,16 +1,14 @@
 """ Top-level module for hessian eigenvec computation """
-from hessian_eigenthings.power_iter import power_iteration,\
- deflated_power_iteration
+from hessian_eigenthings.power_iter import power_iteration, deflated_power_iteration
 from hessian_eigenthings.lanczos import lanczos
-from hessian_eigenthings.hvp_operator import HVPOperator,\
- compute_hessian_eigenthings
+from hessian_eigenthings.hvp_operator import HVPOperator, compute_hessian_eigenthings
 
 __all__ = [
- 'power_iteration',
- 'deflated_power_iteration',
- 'lanczos',
- 'HVPOperator',
- 'compute_hessian_eigenthings'
+ "power_iteration",
+ "deflated_power_iteration",
+ "lanczos",
+ "HVPOperator",
+ "compute_hessian_eigenthings",
 ]
 
-name = 'hessian_eigenthings'
+name = "hessian_eigenthings"

hessian_eigenthings/power_iter.py

@@ -5,6 +5,8 @@
 import numpy as np
 import torch
 
+from hessian_eigenthings.utils import log, progress_bar
+
 
 class Operator:
  """
@@ -59,7 +61,9 @@ def deflated_power_iteration(
  def _deflate(x, val, vec):
  return val * vec.dot(x) * vec
 
- for _ in range(num_eigenthings):
+ log("beginning deflated power iteration")
+ for i in range(num_eigenthings):
+ log("computing eigenvalue/vector %d of %d" % (i + 1, num_eigenthings))
  eigenval, eigenvec = power_iteration(
  current_op,
  power_iter_steps,
@@ -68,6 +72,7 @@ def _deflate(x, val, vec):
  use_gpu=use_gpu,
  init_vec=prev_vec,
  )
+ log("eigenvalue %d: %.4f" % (i + 1, eigenval))
 
  def _new_op_fn(x, op=current_op, val=eigenval, vec=eigenvec):
  return op.apply(x) - _deflate(x, val, vec)
@@ -111,14 +116,14 @@ def power_iteration(
 
  prev_lambda = 0.0
  prev_vec = torch.zeros_like(vec)
- for _ in range(steps):
+ for i in range(steps):
  new_vec = operator.apply(vec) - momentum * prev_vec
  prev_vec = vec / (torch.norm(vec) + 1e-6)
-
  lambda_estimate = vec.dot(new_vec).item()
  diff = lambda_estimate - prev_lambda
  vec = new_vec.detach() / torch.norm(new_vec)
  error = np.abs(diff / lambda_estimate)
+ progress_bar(i, steps, "power iter error: %.4f" % error)
  if error < error_threshold:
  return lambda_estimate, vec
  prev_lambda = lambda_estimate