flennerhag / warpgrad Goto Github PK

class SimpleModel(nn.Module):
    # The second linear layer will be used as warp layer.
    def __init__(self,hidden_dim):
        super(SimpleModel, self).__init__()
        self.linear1 = nn.Linear(2,hidden_dim)
        self.relu1 = nn.ReLU()
        self.linear2 = nn.Linear(hidden_dim,hidden_dim)
        self.relu2 = nn.ReLU()
        self.linear3 = nn.Linear(hidden_dim,2)
    def forward(self,x):
        out = self.linear1(x)
        out = self.relu1(out)
        out = self.linear2(out)
        out = self.relu2(out)
        out = self.linear3(out)
        return out

# It seems like any model that will be given as an argument to the initialization of the `Warp` class should have a function called `init_adaptation`.

updater = warpgrad.DualUpdater(criterion=nn.CrossEntropyLoss)
buffer = warpgrad.ReplayBuffer()

# I will already be going to use inner_optimizer and meta_optimizer for the task adaptation problem.
# What is the connection of these optimization parameters with the algorithm itself?

optimizer_parameters = warpgrad.OptimizerParameters(
                trainable=False,
                default_lr=1e-3,
                default_momentum=0.9)
simple_model = SimpleModel(hidden_dim=100)

model = warpgrad.Warp(simple_model , 
                        [simple_model.linear1,simple_model.relu1,simple_model.relu2,simple_model.linear3],   # List of nn.Modules for adaptation
                        [simple_model.linear2],  # List of nn.Modules for warping
                        updater, 
                        buffer,
                        optimizer_parameters) 

# Should we use the optimizers defined in the Warpgrad library as `meta_opt_class`?
meta_opt = meta_opt_class(model.warp_parameters(), **meta_opt_kwargs)

def meta_step_fn():
    meta_opt.step()
    meta_opt.zero_grad()

for meta_step in range(meta_steps):
    meta_batch = mytaskgenerator.sample()
    for task in meta_batch:
        model.init_adaptation()     # Initialize adaptation on the model side
        # How should task should be defined? It seems like it needs to be a specific class.
        model.register_task(task)   # Register task in replay buffer
        model.collect()             # Turn parameter collection on

        # Is this optimizer should also be from the warpgrad library?
        opt = opt_cls(model.adapt_parameters(), **opt_kwargs)

        for x, y in task:
            loss = criterion(model(x), y)
            loss.backward()
            opt.step()
            opt.zero_grad()

        # Evaluation
        model.eval()
        model.no_collect()
        your_evaluation(task, model)

    # Meta-objective
    model.backward(meta_step_fn)