make eigenvalue float instead of tensor (#9)

Author: noahgolmant

example/main.py

@@ -24,6 +24,8 @@ def extra_args(parser):
  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')
 
 
 def main(args):
@@ -38,7 +40,8 @@ def main(args):
  criterion,
  args.num_eigenthings,
  args.num_steps,
- momentum=args.momentum)
+ momentum=args.momentum,
+ use_gpu=args.cuda)
  print("Eigenvecs:")
  print(eigenvecs)
  print("Eigenvals:")

hessian_eigenthings/power_iter.py

@@ -63,7 +63,7 @@ def _deflate(x, val, vec):
  def _new_op_fn(x, op=current_op, val=eigenval, vec=eigenvec):
  return op.apply(x) - _deflate(x, val, vec)
  current_op = LambdaOperator(_new_op_fn, operator.size)
- eigenvals.append(eigenval.item())
+ eigenvals.append(eigenval)
  eigenvecs.append(eigenvec.cpu())
 
  return eigenvals, eigenvecs
@@ -88,7 +88,7 @@ def power_iteration(operator, steps=20, error_threshold=1e-4,
  new_vec = operator.apply(vec) - momentum * prev_vec
  prev_vec = vec / torch.norm(vec)
 
- lambda_estimate = vec.dot(new_vec)
+ 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)