代码功能
模型结构: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())