元学习的简单示例

代码功能

模型结构：SimpleModel是一个简单的两层全连接神经网络。
元学习过程：在maml_train函数中，每个任务由支持集和查询集组成。模型先在支持集上进行训练，然后在查询集上进行评估，更新元模型参数。
任务生成：通过create_task_data函数生成随机任务数据，用于模拟不同的学习任务。
元训练和微调：在元训练后，代码展示了如何在新任务上进行模型微调和测试。
这个简单示例展示了如何使用元学习方法（MAML）在不同任务之间共享学习经验，并快速适应新任务。

代码

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset

# 构建一个简单的全连接神经网络作为基础学习器
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc1 = nn.Linear(2, 64)
        self.fc2 = nn.Linear(64, 64)
        self.fc3 = nn.Linear(64, 2)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x

# 创建元学习过程
def maml_train(model, meta_optimizer, tasks, n_inner_steps=1, inner_lr=0.01):
    criterion = nn.CrossEntropyLoss()
    
    # 遍历多个任务
    for task in tasks:
        # 模拟支持集和查询集
        support_data, support_labels, query_data, query_labels = task
        
        # 初始化模型参数，用于内循环训练
        inner_model = SimpleModel()
        inner_model.load_state_dict(model.state_dict())
        inner_optimizer = optim.SGD(inner_model.parameters(), lr=inner_lr)
        
        # 在支持集上进行内循环训练
        for _ in range(n_inner_steps):
            pred_support = inner_model(support_data)
            loss_support = criterion(pred_support, support_labels)
            inner_optimizer.zero_grad()
            loss_support.backward()
            inner_optimizer.step()
        
        # 在查询集上评估
        pred_query = inner_model(query_data)
        loss_query = criterion(pred_query, query_labels)
        
        # 计算梯度并更新元模型
        meta_optimizer.zero_grad()
        loss_query.backward()
        meta_optimizer.step()

# 生成一些简单的任务数据
def create_task_data():
    # 随机生成支持集和查询集
    support_data = torch.randn(10, 2)
    support_labels = torch.randint(0, 2, (10,))
    query_data = torch.randn(10, 2)
    query_labels = torch.randint(0, 2, (10,))
    return support_data, support_labels, query_data, query_labels

# 创建多个任务
tasks = [create_task_data() for _ in range(5)]

# 初始化模型和元优化器
model = SimpleModel()
meta_optimizer = optim.Adam(model.parameters(), lr=0.001)

# 进行元训练
maml_train(model, meta_optimizer, tasks)

# 测试新的任务
new_task = create_task_data()
support_data, support_labels, query_data, query_labels = new_task

# 进行模型微调（内循环）
inner_model = SimpleModel()
inner_model.load_state_dict(model.state_dict())
inner_optimizer = optim.SGD(inner_model.parameters(), lr=0.01)
criterion = nn.CrossEntropyLoss()

# 使用支持集进行一次更新
pred_support = inner_model(support_data)
loss_support = criterion(pred_support, support_labels)
inner_optimizer.zero_grad()
loss_support.backward()
inner_optimizer.step()

# 在查询集上测试
pred_query = inner_model(query_data)
print("预测结果:", pred_query.argmax(dim=1).numpy())
print("真实标签:", query_labels.numpy())