Add cuda support

Author: ikostrikov2

README.md

@@ -12,6 +12,7 @@ python main.py
 - [x] Initial implementation
 - [x] Toy data
 - [x] LSTM updates
+- [ ] Refactor, find a better way to organize the modules
 - [ ] Compare with standard optimizers
 - [ ] Real data
 - [ ] More difficult models

main.py

@@ -24,61 +24,84 @@
  help='number of epoch (default: 100)')
 parser.add_argument('--hidden_size', type=int, default=10, metavar='N',
  help='hidden size of the meta optimizer (default: 10)')
+parser.add_argument('--no-cuda', action='store_true', default=False,
+ help='enables CUDA training')
 args = parser.parse_args()
+args.cuda = not args.no_cuda and torch.cuda.is_available()
 
 assert args.optimizer_steps % args.truncated_bptt_step == 0
 
-# Create a meta optimizer that wraps a model into a meta model
-# to keep track of the meta updates.
-meta_optimizer = MetaOptimizer(MetaModel(Model()), args.hidden_size)
-optimizer = optim.Adam(meta_optimizer.parameters(), lr=1e-3)
-loss_fn = lambda f_x, y: (f_x - y).pow(2).mean()
 
-for epoch in range(args.max_epoch):
- decrease_in_loss = 0.0
- for i in range(args.updates_per_epoch):
-
- # Sample a new model
- model = Model()
-
- x, y = get_batch(args.batch_size)
- x, y = Variable(x), Variable(y)
-
- # Compute initial loss of the model
- f_x = model(x)
- initial_loss = loss_fn(f_x, y)
-
- for k in range(args.optimizer_steps // args.truncated_bptt_step):
- # Keep states for truncated BPTT
- meta_optimizer.reset_lstm(keep_states=k > 0, model=model)
-
- loss_sum = 0
- for j in range(args.truncated_bptt_step):
- x, y = get_batch(args.batch_size)
- x, y = Variable(x), Variable(y)
-
- # First we need to compute the gradients of the model
- f_x = model(x)
- loss = loss_fn(f_x, y)
- model.zero_grad()
- loss.backward()
-
- # Perfom a meta update using gradients from model
- # and return the current meta model saved in the optimizer
- meta_model = meta_optimizer.meta_update(model)
-
- # Compute a loss for a step the meta optimizer
- f_x = meta_model(x)
- loss = loss_fn(f_x, y)
- loss_sum += loss
-
- # Update the parameters of the meta optimizer
- meta_optimizer.zero_grad()
- loss_sum.backward()
- optimizer.step()
-
- # Compute relative decrease in the loss function w.r.t initial value
- decrease_in_loss += loss.data[0] / initial_loss.data[0]
-
- print("Epoch: {}, average final/initial loss ratio: {}".format(epoch,
- decrease_in_loss / args.updates_per_epoch))
+def main():
+ # Create a meta optimizer that wraps a model into a meta model
+ # to keep track of the meta updates.
+ meta_model = Model()
+ if args.cuda:
+ meta_model.cuda()
+
+ meta_optimizer = MetaOptimizer(MetaModel(meta_model), args.hidden_size)
+ if args.cuda:
+ meta_optimizer.cuda()
+
+ optimizer = optim.Adam(meta_optimizer.parameters(), lr=1e-3)
+ loss_fn = lambda f_x, y: (f_x - y).pow(2).mean()
+
+ for epoch in range(args.max_epoch):
+ decrease_in_loss = 0.0
+ for i in range(args.updates_per_epoch):
+
+ # Sample a new model
+ model = Model()
+ if args.cuda:
+ model.cuda()
+
+ x, y = get_batch(args.batch_size)
+ x, y = Variable(x), Variable(y)
+ if args.cuda:
+ x, y = x.cuda(), y.cuda()
+
+ # Compute initial loss of the model
+ f_x = model(x)
+ initial_loss = loss_fn(f_x, y)
+
+ for k in range(args.optimizer_steps // args.truncated_bptt_step):
+ # Keep states for truncated BPTT
+ meta_optimizer.reset_lstm(
+ keep_states=k > 0, model=model, use_cuda=args.cuda)
+
+ loss_sum = 0
+ for j in range(args.truncated_bptt_step):
+ x, y = get_batch(args.batch_size)
+ x, y = Variable(x), Variable(y)
+ if args.cuda:
+ x, y = x.cuda(), y.cuda()
+
+ # First we need to compute the gradients of the model
+ f_x = model(x)
+ loss = loss_fn(f_x, y)
+ model.zero_grad()
+ loss.backward()
+
+ # Perfom a meta update using gradients from model
+ # and return the current meta model saved in the optimizer
+ meta_model = meta_optimizer.meta_update(model)
+
+ # Compute a loss for a step the meta optimizer
+ f_x = meta_model(x)
+ loss = loss_fn(f_x, y)
+ loss_sum += loss
+
+ # Update the parameters of the meta optimizer
+ meta_optimizer.zero_grad()
+ loss_sum.backward()
+ optimizer.step()
+
+ # Compute relative decrease in the loss function w.r.t initial
+ # value
+ decrease_in_loss += loss.data[0] / initial_loss.data[0]
+
+ print("Epoch: {}, average final/initial loss ratio: {}".format(epoch,
+ decrease_in_loss / args.updates_per_epoch))
+
+if __name__ == "__main__":
+ main()

meta_optimizer.py

@@ -24,7 +24,7 @@ def __init__(self, model, hidden_size):
  self.linear2.weight.data.mul_(0.1)
  self.linear2.bias.data.fill_(0.0)
 
- def reset_lstm(self, keep_states=False, model=None):
+ def reset_lstm(self, keep_states=False, model=None, use_cuda=False):
  self.meta_model.reset()
  self.meta_model.copy_params_from(model)
 
@@ -34,6 +34,8 @@ def reset_lstm(self, keep_states=False, model=None):
  else:
  self.hx = Variable(torch.zeros(1, self.hidden_size))
  self.cx = Variable(torch.zeros(1, self.hidden_size))
+ if use_cuda:
+ self.hx, self.cx = self.hx.cuda(), self.cx.cuda()
 
  def forward(self, inputs):
  initial_size = inputs.size()