Initial commit

Author: ikostrikov2

data.py

@@ -0,0 +1,8 @@
+import torch
+
+
+def get_batch(batch_size):
+ x = torch.randn(batch_size, 10)
+ x = x - 2 * x.pow(2)
+ y = x.sum(1)
+ return x, y

main.py

@@ -0,0 +1,78 @@
+import argparse
+import operator
+import sys
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from data import get_batch
+from meta_optimizer import MetaOptimizer
+from model import MetaModel, Model
+from torch.autograd import Variable
+
+parser = argparse.ArgumentParser(description='PyTorch REINFORCE example')
+parser.add_argument('--batch_size', type=int, default=16, metavar='N',
+ help='batch size (default: 16)')
+parser.add_argument('--optimizer_steps', type=int, default=10, metavar='N',
+ help='number of meta optimizer steps (default: 10)')
+parser.add_argument('--updates_per_epoch', type=int, default=100, metavar='N',
+ help='updates per epoch (default: 100)')
+parser.add_argument('--max_epoch', type=int, default=100, metavar='N',
+ 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)')
+args = parser.parse_args()
+
+meta_optimizer = MetaOptimizer(args.hidden_size)
+optimizer = optim.Adam(meta_optimizer.parameters(), lr=1e-3)
+
+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()
+
+ # Create a helper class
+ meta_model = MetaModel()
+ meta_model.copy_params_from(model)
+
+ # Reset lstm values of the meta optimizer
+ meta_optimizer.reset_lstm()
+
+ 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 = (f_x - y).pow(2).mean()
+ loss_sum = 0
+ for j in range(args.optimizer_steps):
+ 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 = (f_x - y).pow(2).mean()
+ model.zero_grad()
+ loss.backward()
+
+ # Perfom a meta update
+ meta_optimizer.meta_update(meta_model, model)
+
+ # Compute a loss for a step the meta optimizer
+ f_x = meta_model(x)
+ loss = (f_x - y).pow(2).mean()
+ loss_sum += loss
+
+ # Compute relative decrease in the loss function w.r.t initial value
+ decrease_in_loss += loss.data[0] / initial_loss.data[0]
+
+ # Update the parameters of the meta optimizer
+ meta_optimizer.zero_grad()
+ loss_sum.backward()
+ optimizer.step()
+
+ print("Epoch: {}, average final/initial loss ratio: {}".format(epoch, decrease_in_loss / args.updates_per_epoch))

meta_optimizer.py

@@ -0,0 +1,86 @@
+from functools import reduce
+from operator import mul
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from torch.autograd import Variable
+
+
+class MetaOptimizer(nn.Module):
+
+ def __init__(self, hidden_size):
+ super(MetaOptimizer, self).__init__()
+ self.hidden_size = hidden_size
+
+ self.linear1 = nn.Linear(1, hidden_size)
+
+ self.lstm = nn.LSTMCell(hidden_size, hidden_size)
+
+ self.linear2 = nn.Linear(hidden_size, 1)
+ self.linear2.weight.data.mul_(0.1)
+ self.linear2.bias.data.fill_(0.0)
+
+ self.reset_lstm()
+
+ def reset_lstm(self):
+ self.hx = Variable(torch.zeros(1, self.hidden_size))
+ self.cx = Variable(torch.zeros(1, self.hidden_size))
+
+ def forward(self, inputs):
+ initial_size = inputs.size()
+ x = inputs.view(-1, 1)
+ x = F.tanh(self.linear1(x))
+
+ if x.size(0) != self.hx.size(0):
+ self.hx = self.hx.expand(x.size(0), self.hx.size(1))
+ self.cx = self.hx.expand(x.size(0), self.cx.size(1))
+
+ self.hx, self.cx = self.lstm(x, (self.hx, self.cx))
+ x = self.hx
+
+ x = self.linear2(x)
+ x = x.view(*initial_size)
+ return x
+
+ def meta_update(self, meta_model, model_with_grads):
+ # First we need to create a flat version of parameters and gradients
+ weight_shapes = []
+ bias_shapes = []
+
+ params = []
+ grads = []
+
+ for module in meta_model.children():
+ weight_shapes.append(list(module._parameters['weight'].size()))
+ bias_shapes.append(list(module._parameters['bias'].size()))
+
+ params.append(module._parameters['weight'].view(-1))
+ params.append(module._parameters['bias'].view(-1))
+
+ for module in model_with_grads.children():
+ grads.append(module._parameters['weight'].grad.view(-1))
+ grads.append(module._parameters['bias'].grad.view(-1))
+
+ flat_params = torch.cat(params)
+ flat_grads = torch.cat(grads)
+
+ # Meta update itself
+ flat_params = flat_params + self(flat_grads)
+
+ # Restore original shapes
+ offset = 0
+ for i, module in enumerate(meta_model.children()):
+ weight_flat_size = reduce(mul, weight_shapes[i], 1)
+ bias_flat_size = reduce(mul, bias_shapes[i], 1)
+
+ module._parameters['weight'] = flat_params[
+ offset:offset + weight_flat_size].view(*weight_shapes[i])
+ module._parameters['bias'] = flat_params[
+ offset + weight_flat_size:offset + weight_flat_size + bias_flat_size].view(*bias_shapes[i])
+
+ offset += weight_flat_size + bias_flat_size
+
+ # Finally, copy values from the meta model to the normal one.
+ meta_model.copy_params_to(model_with_grads)

model.py

@@ -0,0 +1,40 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from torch.autograd import Variable
+
+
+class Model(nn.Module):
+
+ def __init__(self):
+ super(Model, self).__init__()
+ self.linear1 = nn.Linear(10, 16)
+ self.linear2 = nn.Linear(16, 1)
+
+ def forward(self, inputs):
+ x = F.tanh(self.linear1(inputs))
+ x = self.linear2(x)
+ return x
+
+# A helper class that keeps track of meta updates
+# It's done by replacing parameters with variables and applying updates to
+# them.
+
+
+class MetaModel(Model):
+
+ def reset(self):
+ for module in self.children():
+ module._parameters['weight'] = Variable(
+ module._parameters['weight'].data)
+ module._parameters['bias'] = Variable(
+ module._parameters['bias'].data)
+
+ def copy_params_from(self, model):
+ for modelA, modelB in zip(self.parameters(), model.parameters()):
+ modelA.data.copy_(modelB.data)
+
+ def copy_params_to(self, model):
+ for modelA, modelB in zip(self.parameters(), model.parameters()):
+ modelB.data.copy_(modelA.data)