如何自动生成100个含随机样本的字典式PyTorch数据集？

问题背景 #

现有如下形式的PyTorch数据集：

import torch

D1 = {'query': torch.tensor([(1, 2)], dtype=torch.float32),
      'support': torch.tensor([(2, 4), (3, 1)], dtype=torch.float32)}  # (x, y) pairs for query (Q1) and support (S1) set. 

D2 = {'query': torch.tensor([(4, 1)], dtype=torch.float32), 'support': torch.tensor([(5, 3), (6, 0)], dtype=torch.float32)}

每个数据集包含support（训练用）和query（测试/验证用）两个子集，现需自动化生成100个此类数据集，并支持自定义两个子集的随机样本数量。

实现方案 #

可以通过封装一个生成函数批量创建数据集，核心利用PyTorch的随机张量生成能力，自定义样本数量和数据范围：

import torch

def generate_datasets(num_datasets=100, support_size=2, query_size=1, data_min=0., data_max=10.):
    datasets = []
    for _ in range(num_datasets):
        # 生成support集：shape为(support_size, 2)，对应(x,y)对
        support = torch.rand((support_size, 2), dtype=torch.float32) * (data_max - data_min) + data_min
        # 生成query集：shape为(query_size, 2)
        query = torch.rand((query_size, 2), dtype=torch.float32) * (data_max - data_min) + data_min
        datasets.append({'support': support, 'query': query})
    return datasets

# 生成100个数据集，每个support含5个样本、query含3个样本（可按需调整参数）
dataset_list = generate_datasets(num_datasets=100, support_size=5, query_size=3)

# 验证第一个数据集的结构
print("第一个数据集的support形状:", dataset_list[0]['support'].shape)
print("第一个数据集的query形状:", dataset_list[0]['query'].shape)

关键说明 #

• 参数自定义：num_datasets控制生成的数据集总数，support_size和query_size分别设置两个子集的样本数量，data_min和data_max可调整(x,y)数据的取值范围。
• 可复现性：若需要固定随机结果，可在函数开头添加torch.manual_seed(42)（替换为指定种子值）。
• 格式兼容：生成的数据集结构与原D1、D2完全一致，可直接对接后续训练或测试流程。

内容的提问来源于stack exchange，提问作者S.EB