如何为TensorFlow Datasets编写每类保留5张图的少样本过滤器？

我尝试创建带计数属性的对象来跟踪各类样本数量，filter()能接收输入并执行代码，但无法正确统计各类样本数。代码示例如下：

import tensorflow_datasets as tfds

class few_shot_filter:
    def __init__(self):
        self.count = {}

    def filter(self, x):
        y = x["label"]
        if not y in self.count:
            self.count[y] = 0
        self.count[y] = self.count[y] + 1
        return self.count[y] <= 5

ds = tfds.load("cifar10", supervised=False)
fs_filter=few_shot_filter()
print([y["label"] for y in ds.filter(fs_filter.filter).as_numpy_iterator()])

有趣的是，传递对象并访问其属性本身似乎没有问题，如下方代码可正常运行：

import tensorflow_datasets as tfds

class class_filter:
    def __init__(self):
        self.target_class = 5

    def filter(self, x):
        y = x["label"]
        return y == self.target_class

ds = tfds.load("cifar10", supervised=False)
cl_filter=class_filter()
print([y["label"] for y in ds.filter(cl_filter.filter).as_numpy_iterator()])

是不是filter是并行执行的，类似数组索引的方式？

原因分析 #

没错，tf.data的操作默认是并行执行的（多线程读取/处理样本）。你的第一个示例中，每个工作线程都会复制few_shot_filter实例，每个线程拥有独立的self.count字典，导致计数被多份副本分散统计，最终结果错误。

而第二个示例中的target_class是固定值，不需要修改，每个线程的副本都能正确判断标签是否匹配，因此可以正常运行。

解决方案 #

推荐使用按标签分组后取前N个样本的方式，这种方法更可靠，避免并行计数的问题：

import tensorflow_datasets as tfds
from collections import Counter

# 加载cifar10训练集
ds = tfds.load("cifar10", split="train", supervised=False)

# 定义每组保留前5个样本的函数
def keep_top_n_samples(ds):
    return ds.take(5)

# 按标签分组，对每组应用保留函数
ds_few_shot = ds.group_by_window(
    key_func=lambda x: x["label"],  # 分组依据：样本标签
    reduce_func=lambda key, ds: keep_top_n_samples(ds),
    window_size=10000  # 窗口值设为大于每类样本数，确保覆盖所有同标签样本
)

# 验证结果
labels = [y["label"] for y in ds_few_shot.as_numpy_iterator()]
print(Counter(labels))  # 输出：Counter({0:5, 1:5, ..., 9:5})

备选方案：使用tf.Variable共享计数（不推荐，存在并行风险） #

如果一定要用filter实现，可以用tf.Variable存储计数（图中共享变量），并通过原子操作更新计数，但仍可能因并行执行顺序问题导致结果偏差：

import tensorflow as tf
import tensorflow_datasets as tfds
from collections import Counter

class FewShotFilter:
    def __init__(self, num_classes=10):
        # 用tf.Variable存储每类计数，确保在图中共享
        self.count = tf.Variable(tf.zeros(num_classes, dtype=tf.int32))
    
    @tf.autograph.experimental.do_not_convert
    def filter_fn(self, x):
        label = x["label"]
        # 原子更新计数：获取当前值→加1→更新变量
        current_count = tf.gather(self.count, label)
        new_count = current_count + 1
        self.count.assign(tf.tensor_scatter_nd_update(
            self.count,
            indices=[[label]],
            updates=[new_count]
        ))
        return new_count <= 5

ds = tfds.load("cifar10", split="train", supervised=False)
fs_filter = FewShotFilter()
ds_few_shot = ds.filter(fs_filter.filter_fn)

# 验证
labels = [y["label"] for y in ds_few_shot.as_numpy_iterator()]
print(Counter(labels))

内容的提问来源于stack exchange，提问作者frebo2847