狗都能看懂的Pytorch MAML代码详解

源码（觉得有用请点star，这对我很重要~）

maml概念

首先，我们需要说明的是maml不同于常见的训练方式。以猫狗分类和resnet作为例子，我们将猫狗分类定义为一个task，正常训练一个猫狗分类器，只需要输入猫和狗的图片去训练就好了。所以我们的一个batch中就会有多张猫或者狗的图片，这样训练出来的模型虽说可以预测这张图片是猫还是狗，但要想这个分类器有泛化性，就需要大量猫或狗的图像，而标注大量的数据是要成本的。

现在我们假设一个场景，我们没有这么多猫狗分类的数据，但我们有其他task的数据。我们需要用少量的图像来训练一个强泛化性的模型。maml的训练方式允许我们用大量别的task的数据来得到一个初始化权重，这个初始化权重具有非常好的鲁棒性，仅用少量数据训练加上或者maml训练的初始化权重就可以达到和正常训练方式的准确率。

为什么maml能做到这样的效果，请读者移步MAML原理讲解和代码实现。

maml以task为单位，多个task组成一个batch，为了和正常训练方式区别开来，我们就成为meta-batch。以omniglot为例，如下图所示：

每个task之间都互相独立，都是不同的分类任务。

数据读取

这里为大家实现了个MAML数据读取的基类，用户只需要实现get_file_list和get_one_task_data两个函数即可。

class MAMLDataset(Dataset):

    def __init__(self, data_path, batch_size, n_way=10, k_shot=2, q_query=1):

        self.file_list = self.get_file_list(data_path)
        self.batch_size = batch_size
        self.n_way = n_way
        self.k_shot = k_shot
        self.q_query = q_query

    def get_file_list(self, data_path):
        raise NotImplementedError('get_file_list function not implemented!')

    def get_one_task_data(self):
        raise NotImplementedError('get_one_task_data function not implemented!')

    def __len__(self):
        return len(self.file_list) // self.batch_size

    def __getitem__(self, index):
        return self.get_one_task_data()

还是以omniglot为例，实现特殊数据集的子类数据读取的方法。

get_file_list

此函数要求得到一个所有task文件目录的list。比如一个总的文件夹中，下面有很多不同的task，这里因为omniglot数据命名比较统一，所以实现比较简单。

get_one_task_data

此函数要求返回一个task的数据，包括训练集和验证集，以下面代码为例，每次调用get_one_task_data时，返回一个n_way=5分类的task，其中训练集图像和标签的数量各为k_shot=1张，验证集图像和标签的数量各为q_query=1张。

class OmniglotDataset(MAMLDataset):
    def get_file_list(self, data_path):
        """
        Get all fonts list.
        Args:
            data_path: Omniglot Data path

        Returns: fonts list

        """
        return [f for f in glob.glob(data_path + "**/character*", recursive=True)]

    def get_one_task_data(self):
        """
        Get ones task maml data, include one batch support images and labels, one batch query images and labels.
        Returns: support_data, query_data

        """
        img_dirs = random.sample(self.file_list, self.n_way)
        support_data = []
        query_data = []

        support_image = []
        support_label = []
        query_image = []
        query_label = []

        for label, img_dir in enumerate(img_dirs):
            img_list = [f for f in glob.glob(img_dir + "**/*.png", recursive=True)]
            images = random.sample(img_list, self.k_shot + self.q_query)

            # Read support set
            for img_path in images[:self.k_shot]:
                image = Image.open(img_path)
                image = np.array(image)
                image = np.expand_dims(image / 255., axis=0)
                support_data.append((image, label))

            # Read query set
            for img_path in images[self.k_shot:]:
                image = Image.open(img_path)
                image = np.array(image)
                image = np.expand_dims(image / 255., axis=0)
                query_data.append((image, label))

        # shuffle support set
        random.shuffle(support_data)
        for data in support_data:
            support_image.append(data[0])
            support_label.append(data[1])

        # shuffle query set
        random.shuffle(query_data)
        for data in query_data:
            query_image.append(data[0])
            query_label.append(data[1])

        return np.array(support_image), np.array(support_label), np.array(query_image), np.array(query_label)

在调用Dataset的时候再使用torch的Dataloader包一下就好了，里面batch_size参数为任务的数量。相当于每训练1个step就要训练完这么多个task。

train_dataset = OmniglotDataset(args.train_data_dir, args.task_num,
                                n_way=args.n_way, k_shot=args.k_shot, q_query=args.q_query)
val_dataset = OmniglotDataset(args.val_data_dir, args.val_task_num,
                              n_way=args.n_way, k_shot=args.k_shot, q_query=args.q_query)
train_loader = DataLoader(train_dataset, batch_size=args.task_num, shuffle=True, num_workers=args.num_workers)
val_loader = DataLoader(val_dataset, batch_size=args.val_task_num, shuffle=False, num_workers=args.num_workers)

模型训练

代码如下：

def maml_train(model, support_images, support_labels, query_images, query_labels, inner_step, args, optimizer, is_train=True):
    """
    Train the model using MAML method.
    Args:
        model: Any model
        support_images: several task support images
        support_labels: several  support labels
        query_images: several query images
        query_labels: several query labels
        inner_step: support data training step
        args: ArgumentParser
        optimizer: optimizer
        is_train: whether train

    Returns: meta loss, meta accuracy

    """
    meta_loss = []
    meta_acc = []

    for support_image, support_label, query_image, query_label in zip(support_images, support_labels, query_images, query_labels):

        fast_weights = collections.OrderedDict(model.named_parameters())
        for _ in range(inner_step):
            # Update weight
            support_logit = model.functional_forward(support_image, fast_weights)
            support_loss = nn.CrossEntropyLoss().cuda()(support_logit, support_label)
            grads = torch.autograd.grad(support_loss, fast_weights.values(), create_graph=True)
            fast_weights = collections.OrderedDict((name, param - args.inner_lr * grads)
                                                   for ((name, param), grads) in zip(fast_weights.items(), grads))

        # Use trained weight to get query loss
        query_logit = model.functional_forward(query_image, fast_weights)
        query_prediction = torch.max(query_logit, dim=1)[1]

        query_loss = nn.CrossEntropyLoss().cuda()(query_logit, query_label)
        query_acc = torch.eq(query_label, query_prediction).sum() / len(query_label)

        meta_loss.append(query_loss)
        meta_acc.append(query_acc.data.cpu().numpy())

    # Zero the gradient
    optimizer.zero_grad()
    meta_loss = torch.stack(meta_loss).mean()
    meta_acc = np.mean(meta_acc)

    if is_train:
        meta_loss.backward()
        optimizer.step()

    return meta_loss, meta_acc

support_images, support_labels, query_images, query_labels传入的都是以task为单位的，所以要用一个for循环来进行拆包，注意support_data和query_data数据集来源必须得一致，不能一个数据A task，另一个属于B task。

拆包完之后，首先进行训练集的训练，我们要注意，此时的训练是不能改动到模型权重，但我们又需要知道它的训练方向，所以我们需要copy出来一个权重，让它执行训练，用这个得到的权重对query_data执行前向传播，以此得到的loss再进行反向传播优化。这个过程很绕，建议多读几遍源码就懂了。

模型定义

class Classifier(nn.Module):
    def __init__(self, in_ch, n_way):
        super(Classifier, self).__init__()
        self.conv1 = ConvBlock(in_ch, 64)
        self.conv2 = ConvBlock(64, 64)
        self.conv3 = ConvBlock(64, 64)
        self.conv4 = ConvBlock(64, 64)
        self.logits = nn.Linear(64, n_way)

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.conv3(x)
        x = self.conv4(x)
        x = x.view(x.shape[0], -1)
        x = self.logits(x)

        return x

    def functional_forward(self, x, params):
        x = ConvBlockFunction(x, params[f'conv1.conv2d.weight'], params[f'conv1.conv2d.bias'],
                              params.get(f'conv1.bn.weight'), params.get(f'conv1.bn.bias'))
        x = ConvBlockFunction(x, params[f'conv2.conv2d.weight'], params[f'conv2.conv2d.bias'],
                              params.get(f'conv2.bn.weight'), params.get(f'conv2.bn.bias'))
        x = ConvBlockFunction(x, params[f'conv3.conv2d.weight'], params[f'conv3.conv2d.bias'],
                              params.get(f'conv3.bn.weight'), params.get(f'conv3.bn.bias'))
        x = ConvBlockFunction(x, params[f'conv4.conv2d.weight'], params[f'conv4.conv2d.bias'],
                              params.get(f'conv4.bn.weight'), params.get(f'conv4.bn.bias'))

        x = x.view(x.shape[0], -1)
        x = F.linear(x, params['logits.weight'], params['logits.bias'])

        return x

模型定义比较简单，maml思想主要是个训练方式，和模型本身无关。但我们在刚刚模型训练的时候有一些特殊操作，所以要定义一个functional_forward，这个函数要求实现和模型一样结构的网络，同时参数输入为：1、图像 2、权重。这样就可以保证得到了loss，但模型权重没有被修改。