MAML_Example.py

import torch
import torch.nn as nn
import torch.optim as optim

# Define the model architecture
class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.linear = nn.Linear(1, 1)

    def forward(self, x):
        return self.linear(x)

# Define the MAML algorithm
class MAML:
    def __init__(self, model, lr_inner=0.01, lr_outer=0.001):
        self.model = model
        self.lr_inner = lr_inner
        self.lr_outer = lr_outer

    def inner_loop(self, task, x, y):
        task_model = Model()
        task_model.load_state_dict(self.model.state_dict())
        optimizer = optim.SGD(task_model.parameters(), lr=self.lr_inner)
        for i in range(10):
            loss = nn.MSELoss()(task_model(x), y)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
        return task_model

    def outer_loop(self, tasks):
        optimizer = optim.Adam(self.model.parameters(), lr=self.lr_outer)
        for task in tasks:
            x, y = task
            task_model = self.inner_loop(task, x, y)
            loss = nn.MSELoss()(task_model(x), y)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

# Define the tasks
tasks = []
for i in range(10):
    x = torch.randn(10, 1)
    y = x * 2 + torch.randn(10, 1) * 0.1
    tasks.append((x, y))

# Initialize the model and the MAML algorithm
model = Model()
maml = MAML(model)

# Train the model using MAML
maml.outer_loop(tasks)