np.random.choice是否有一个Tensorflow等价物，用于从离散集合中随机采样？

Question

我正在尝试在TensorFlow中实现一种类似于和的图像增强策略。从RandAugment论文中，以下代码显示了如何随机选择N增强应用于图像

transforms = [’Identity’, ’AutoContrast’, ’Equalize’, ’Rotate’, ’Solarize’,
              ’Color’,’Posterize’, ’Contrast’, ’Brightness’, ’Sharpness’,
              ’ShearX’, ’ShearY’,’TranslateX’, ’TranslateY’]

def randaugment(N, M):
"""Generate a set of distortions.
Args:
N: Number of augmentation transformations to apply
sequentially.
M: Magnitude for all the transformations.
"""
sampled_ops = np.random.choice(transforms, N)
return [(op, M) for op in sampled_ops]

然而，我希望在TensorFlow中对每一批图像执行此操作，理想情况下尽可能高效。看起来像

transform_names = ['Identity', 'Brightness', 'Colour', 'Contrast', 'Equalise', 'Rotate',
                   'Sharpness', 'ShearX', 'ShearY', 'TranslateX', 'TranslateY']

transforms = {'Identity':identity, 'Brightness':brightness, 'Colour':colour,
              'Contrast':contrast, 'Equalise':equalise, 'Rotate':rotate,
              'Sharpness':sharpness, 'ShearX':shear_x, 'ShearY':shear_y,
              'TranslateX':translate_x, 'TranslateY':translate_y}

def brightness(image, M):
    M = tf.math.minimum(M, 0.95)
    M = tf.math.maximum(M, 0.05)
    B = M - 1
    image = tf.image.adjust_brightness(image, delta=B)
    image = tf.clip_by_value(image, clip_value_min=0, clip_value_max=1)

    return image

def augment(image):
    N = 3
    M = tf.random.uniform(minval=0, maxval=1, shape=[])
    sampled_ops = np.random.choice(transform_names, N)

    for op in sampled_ops:
        image = transforms[op](image, M)

    return image

x = tf.data.Dataset.from_tensor_slices(x)
x = x.batch(batch_size)
x_a = x.map(augment)

其中，x是图像数据集，augment是随机采样N增强以应用于每个图像的增强函数。我添加了亮度函数来说明各个增强函数的组成。从我收集的数据来看，任何NumPy函数似乎只在整个数据集中调用一次，这意味着每个图像的采样增强都是相同的

我怎样才能编写这样的代码，使得每个批次的单个增强都是独立随机抽样的