wine_QNNimprove
import sys
import os
import pennylane as qml
from pennylane import numpy as np
import tensorflow as tf
from sklearn.datasets import load_wine
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler, OneHotEncoder
from sklearn.decomposition import PCA
import matplotlib
matplotlib.use('TkAgg')
import matplotlib.pyplot as plt
from tensorflow.keras.layers import Dense
from importance_sampling.training import ConstantTimeImportanceTraining # 确保该库已安装

plt.rcParams['font.sans-serif'] = ['SimHei'] # 使用黑体
plt.rcParams['axes.unicode_minus'] = False # 解决负号显示问题

# 重定向 stdout 以捕获打印输出
class Tee(object):
def __init__(self, *files):
self.files = files
def write(self, obj):
for f in self.files:
f.write(obj)
f.flush()
def flush(self):
for f in self.files:
f.flush()

output_file = open('QNNImprove', 'w', encoding='utf-8')
original_stdout = sys.stdout
sys.stdout = Tee(sys.stdout, output_file)

def save_plot(fig, filename):
fig.savefig(filename)
plt.close(fig)

def plot_history(history):
# 绘制训练和验证准确率
fig, ax = plt.subplots(figsize=(10, 5))
ax.plot(history.history['accuracy'])
ax.plot(history.history['val_accuracy'])
ax.set_title('模型准确率')
ax.set_ylabel('准确率')
ax.set_xlabel('轮次')
ax.legend(['训练', '验证'], loc='upper left')
save_plot(fig, 'accuracy_plot.png')

# 绘制训练和验证损失
fig, ax = plt.subplots(figsize=(10, 5))
ax.plot(history.history['loss'])
ax.plot(history.history['val_loss'])
ax.set_title('模型损失')
ax.set_ylabel('损失')
ax.set_xlabel('轮次')
ax.legend(['训练', '验证'], loc='upper left')
save_plot(fig, 'loss_plot.png')

def calculate_accuracy_fluctuation(history):
val_accuracies = history.history['val_accuracy']
mean_accuracy = np.mean(val_accuracies)
fluctuation = np.sum(np.abs(val_accuracies - mean_accuracy))
return mean_accuracy, fluctuation

# 定义设备，使用 8 个量子比特
n_qubits = 8 # 根据需要的特征数量调整量子比特数
dev = qml.device("default.qubit", wires=n_qubits)

def encode_features(inputs, n_qubits):
"""
使用角度编码将输入特征编码到量子电路中。
将特征映射到量子比特的旋转角度。
"""
# 将特征映射到量子比特
for i in range(n_qubits):
qml.RY(inputs[i], wires=i)

@qml.qnode(dev, interface='tf', diff_method='backprop')
def quantum_neural_network(inputs, rot_weights):
# 编码输入特征到量子比特
encode_features(inputs, n_qubits)

# 应用变分层
for layer in range(rot_weights.shape[0]):
for qubit in range(n_qubits):
rot_angles = rot_weights[layer, qubit, :] # 形状：(3,)
qml.Rot(rot_angles[0], rot_angles[1], rot_angles[2], wires=qubit)
# 修改纠缠方式，使用全连接纠缠
for i in range(n_qubits):
for j in range(i + 1, n_qubits):
qml.CRX(np.pi / 2, wires=[i, j])

# 测量期望值
return [qml.expval(qml.PauliZ(i)) for i in range(n_qubits)]

# 定义用于绘制电路的 QNode
@qml.qnode(dev)
def quantum_neural_network_draw(inputs, rot_weights):
# 编码输入特征到量子比特
encode_features(inputs, n_qubits)

# 应用变分层
for layer in range(rot_weights.shape[0]):
for qubit in range(n_qubits):
rot_angles = rot_weights[layer, qubit, :] # 形状：(3,)
qml.Rot(rot_angles[0], rot_angles[1], rot_angles[2], wires=qubit)
# 修改纠缠方式，使用全连接纠缠
for i in range(n_qubits):
for j in range(i + 1, n_qubits):
qml.CRX(np.pi / 2, wires=[i, j])

# 测量期望值
return [qml.expval(qml.PauliZ(i)) for i in range(n_qubits)]

# 定义量子神经网络层
class QuantumNeuralNetworkLayer(tf.keras.layers.Layer):
def __init__(self, n_qubits, n_layers, **kwargs):
super(QuantumNeuralNetworkLayer, self).__init__(**kwargs)
self.n_qubits = n_qubits
self.n_layers = n_layers

# 定义旋转权重
init_weights = tf.random.uniform(
shape=(n_layers, n_qubits, 3),
minval=0,
maxval=2 * np.pi,
dtype=tf.float32
)
self.rot_weights = tf.Variable(
initial_value=init_weights,
trainable=True,
name='rot_weights'
)

def call(self, inputs):
def apply_qnode(x):
x = tf.cast(x, dtype=tf.float32)
result = quantum_neural_network(x, self.rot_weights)
return tf.cast(result, dtype=tf.float32)

outputs = tf.map_fn(
apply_qnode,
inputs,
dtype=tf.float32
)

return outputs

# 定义量子神经网络模型
class QuantumNeuralNetwork(tf.keras.Model):
def __init__(self, n_qubits, n_layers, num_classes, **kwargs):
super(QuantumNeuralNetwork, self).__init__(**kwargs)
self.n_qubits = n_qubits
self.n_layers = n_layers
self.num_classes = num_classes

# 量子层
self.quantum_layer = QuantumNeuralNetworkLayer(n_qubits, n_layers)

# 经典输出层
self.output_layer = tf.keras.layers.Dense(num_classes, activation='softmax')

def call(self, inputs):
q_output = self.quantum_layer(inputs)
return self.output_layer(q_output)

if __name__ == "__main__":
# **重要性采样始终启用**

# 加载 Wine 数据集
wine = load_wine()
X = wine['data'].astype(np.float32)
y = wine['target'].reshape(-1, 1)

# One-hot 编码目标标签
encoder = OneHotEncoder(sparse_output=False)
y_encoded = encoder.fit_transform(y).astype(np.float32)
num_classes = y_encoded.shape[1]

# 先划分数据集
x_train, x_test, y_train, y_test = train_test_split(
X, y_encoded, test_size=0.3, random_state=42
)

# 特征标准化 (仅在训练集上拟合)
scaler = StandardScaler()
x_train_scaled = scaler.fit_transform(x_train)
x_test_scaled = scaler.transform(x_test)

# 使用 PCA 进行降维 (仅在训练集上拟合)
n_features = n_qubits # 将特征数量降至与量子比特数量一致
pca = PCA(n_components=n_features)
x_train_reduced = pca.fit_transform(x_train_scaled)
x_test_reduced = pca.transform(x_test_scaled)

# 将特征归一化到 [0, π] (基于训练集的最小值和最大值)
x_min = x_train_reduced.min(axis=0)
x_max = x_train_reduced.max(axis=0)
x_train_normalized = (x_train_reduced - x_min) / (x_max - x_min) * np.pi
x_test_normalized = (x_test_reduced - x_min) / (x_max - x_min) * np.pi

x_train = x_train_normalized
x_test = x_test_normalized

learning_rate = 0.01 # 学习率
epochs = 200 # 训练轮次
batch_size = 16
n_layers = 5 # 增加量子层数

model = QuantumNeuralNetwork(
n_qubits=n_qubits,
n_layers=n_layers,
num_classes=num_classes
)

# 使用 Adam 优化器编译模型
model.compile(
optimizer=tf.keras.optimizers.Adam(learning_rate=learning_rate),
loss='categorical_crossentropy',
metrics=['accuracy']
)

def is_quantum_model(model):
for layer in model.layers:
if isinstance(layer, QuantumNeuralNetworkLayer):
return True
return False

if is_quantum_model(model):
print("模型包含量子神经网络（QNN）层。")
else:
print("模型不包含量子神经网络（QNN）层。")

print("\n量子电路图：")
sample_inputs = np.ones(n_qubits, dtype=np.float32) * (np.pi / 2) # 使用 n_qubits 个特征
sample_rot_weights = np.zeros((n_layers, n_qubits, 3), dtype=np.float32) # rot_weights

print(qml.draw(quantum_neural_network_draw)(sample_inputs, sample_rot_weights)) # 使用绘制专用的 QNode

# 早停和学习率调整回调函数
callbacks = [
tf.keras.callbacks.EarlyStopping(monitor='val_loss', patience=10, restore_best_weights=True),
tf.keras.callbacks.ReduceLROnPlateau(monitor='val_loss', factor=0.5, patience=15)
]

# **始终使用重要性采样进行训练**
wrapped_model = ConstantTimeImportanceTraining(model, score="loss")
history = wrapped_model.fit(
x_train, y_train,
batch_size=batch_size,
epochs=epochs,
verbose=1,
validation_data=(x_test, y_test),
callbacks=callbacks
)

score = model.evaluate(x_test, y_test, verbose=0)
print('测试损失:', score[0])
print('测试准确率:', score[1])

mean_accuracy, fluctuation = calculate_accuracy_fluctuation(history)
print('平均验证准确率:', mean_accuracy)
print('准确率波动:', fluctuation)

plot_history(history)

# 关闭输出文件并恢复 stdout
sys.stdout = original_stdout
output_file.close()

print("脚本执行完成。查看 'output.log' 获取打印输出，查看当前目录获取保存的图表。")
上述是wine_QNNimprove.py文件

import numpy as np
from blinker import signal
from tensorflow.python.keras import backend as K
from tensorflow.python.keras.layers import Layer
from tensorflow.python.keras.metrics import categorical_accuracy
import tensorflow as tf
from .layers import GradientNormLayer, LossLayer, MetricLayer

class ModelWrapper(object):
"""ModelWrapper 的目的是包装一个神经网络，并添加一些额外的层，以生成得分、损失和样本权重，以执行重要性采样。"""

def _iterate_batches(self, x, y, batch_size):
bs = batch_size
for s in range(0, len(y), bs):
if isinstance(x, (list, tuple)):
xs = [xi[s:s+bs] for xi in x]
else:
xs = x[s:s+bs]
yield xs, y[s:s+bs]

def evaluate(self, x, y, batch_size=128):
result = np.mean(
np.vstack([
self.evaluate_batch(xi, yi)
for xi, yi in self._iterate_batches(x, y, batch_size)
]),
axis=0
)

signal("is.evaluation").send(result)
return result

def score(self, x, y, batch_size=128):
bs = batch_size
result = np.hstack([
self.score_batch(xi, yi).T
for xi, yi in self._iterate_batches(x, y, batch_size)
]).T

signal("is.score").send(result)
return result

def set_lr(self, lr):
"""设置模型的学习率。"""

try:
K.set_value(
self.optimizer.lr,
lr
)
except AttributeError:
try:
K.set_value(
self.model.optimizer.lr,
lr
)
except AttributeError:
raise NotImplementedError()

try:
K.set_value(
self.small.optimizer.lr,
lr
)
except AttributeError:
pass

def evaluate_batch(self, x, y):
raise NotImplementedError()

def score_batch(self, x, y):
raise NotImplementedError()

def train_batch(self, x, y, w):
raise NotImplementedError()

class OracleWrapper(ModelWrapper):
AVG_LOSS = 0
LOSS = 1
WEIGHTED_LOSS = 2
SCORE = 3
METRIC0 = 4

FUSED_ACTIVATION_WARNING = ("[WARNING]: 最后一层具有融合激活函数，例如 Dense(..., activation=\"sigmoid\")。\n"
"为了自动提取预激活张量，请使用单独的激活层（参见示例）。\n")

def __init__(self, model, reweighting, score="loss", layer=None):
self.reweighting = reweighting
self.layer = self._gnorm_layer(model, layer)
self._augment_model(model, score, reweighting)

def _gnorm_layer(self, model, layer):
# 如果提供了层，则直接使用
if isinstance(layer, Layer):
return layer

# 如果提供了层的索引，则提取该层
if isinstance(layer, int):
return model.layers[layer]

# 对于量子神经网络，可能无法自动获取梯度范数层
# 这里暂时返回 None，表示不使用特定的层
return None

def _augment_model(self, model, score, reweighting):
# 保存原始模型
self.original_model = model

# 创建包装后的模型，不修改原始模型的结构
self.model = model # 使用原始模型

# 保存必要的信息
self.score_method = score
self.reweighting = reweighting

# 提取模型的一些信息
self.loss_fn = model.loss
self.optimizer = model.optimizer.__class__.from_config(model.optimizer.get_config())
self.metrics = model.metrics or []

def evaluate_batch(self, x, y):
# 使用 model.evaluate
outputs = self.model.evaluate(
x=x,
y=y,
verbose=0,
return_dict=True
)
loss = outputs['loss']
metrics = [outputs[m] for m in self.model.metrics_names[1:]]

signal("is.evaluate_batch").send(outputs)

# 返回损失和其他指标
return [loss] + metrics

# 在 model_wrappers.py 文件中
def score_batch(self, x, y):
y_pred = self.model(x, training=False)
loss_fn = self.model.loss

# 如果 loss_fn 是字符串，使用 tf.keras.losses.get 将其转换为损失函数对象
if isinstance(loss_fn, str):
loss_fn = tf.keras.losses.get(loss_fn)
elif isinstance(loss_fn, dict):
loss_fn = tf.keras.losses.get(loss_fn)
elif not callable(loss_fn):
raise ValueError(f"Unsupported loss function type: {type(loss_fn)}")

# 确保损失函数具有 reduction 参数，并设置为 NONE
if isinstance(loss_fn, tf.keras.losses.Loss):
loss_config = loss_fn.get_config()
per_example_loss_fn = loss_fn.__class__(reduction=tf.keras.losses.Reduction.NONE, **loss_config)
per_sample_losses = per_example_loss_fn(y, y_pred)
else:
per_sample_losses = loss_fn(y, y_pred)

if self.score_method == "loss":
scores = per_sample_losses
else:
scores = tf.ones_like(per_sample_losses)

return scores.numpy()

def train_batch(self, x, y, w):
with tf.GradientTape() as tape:
y_pred = self.model(x, training=True)
# 计算每个样本的损失（不进行 reduction）
per_sample_losses = self.model.compiled_loss(
y,
y_pred,
sample_weight=None,
regularization_losses=self.model.losses,
reduction=tf.keras.losses.Reduction.NONE
)
# 计算得分
if self.score_method == "loss":
scores = per_sample_losses
else:
# 实现其他得分方式
scores = tf.ones_like(per_sample_losses)

# 计算样本权重
weights = self.reweighting.weight_layer()([scores, w])

# 计算加权损失
weighted_losses = per_sample_losses * weights
loss = tf.reduce_mean(weighted_losses)

# 获取模型的可训练参数
trainable_vars = self.model.trainable_variables

# 计算梯度
gradients = tape.gradient(loss, trainable_vars)

# 更新模型参数
self.model.optimizer.apply_gradients(zip(gradients, trainable_vars))

# 更新指标
self.model.compiled_metrics.update_state(y, y_pred, sample_weight=weights)
metrics = [metric.result().numpy() for metric in self.model.metrics]

# 返回损失、指标和得分
return loss.numpy(), metrics, scores.numpy()

def __call__(self, x, training=False):
return self.model(x, training=training)
def _tolist(x, acceptable_iterables=(list, tuple)):
if not isinstance(x, acceptable_iterables):
x = [x]
return x

def _get_scoring_layer(score, y_true, y_pred, loss_fn="categorical_crossentropy",
layer=None, model=None):
"""获取一个计算每个 y_true 和 y_pred 对的得分的层"""
assert score in ["loss", "gnorm", "full_gnorm", "acc"]

if score == "loss":
return LossLayer(loss_fn, name='score_layer')([
y_true,
y_pred
])
elif score == "gnorm":
return GradientNormLayer(
layer.output,
loss_fn,
fast=False,
name='score_layer'
)([
y_true,
y_pred
])
elif score == "full_gnorm":
return GradientNormLayer(
model.trainable_weights,
loss_fn,
fast=False,
name='score_layer'
)([
y_true,
y_pred
])
elif score == "acc":
return MetricLayer(categorical_accuracy, name='score_layer')([
y_true,
y_pred
])

class ModelWrapperDecorator(ModelWrapper):
def __init__(self, model_wrapper, implemented_attributes=set()):
self.model_wrapper = model_wrapper
self.implemented_attributes = (
implemented_attributes | set(["model_wrapper"])
)

def __getattribute__(self, name):
_getattr = object.__getattribute__
implemented_attributes = _getattr(self, "implemented_attributes")
if name in implemented_attributes:
return _getattr(self, name)
else:
model_wrapper = _getattr(self, "model_wrapper")
return getattr(model_wrapper, name)
上述是model_wrappers.py文件
import sys
import numpy as np
import logging
from collections.abc import Iterable
import tensorflow as tf
from tensorflow.python.keras.callbacks import BaseLogger, CallbackList, History, ProgbarLogger
from tensorflow.python.keras.layers import Dropout
#from tensorflow.python.keras.utils import Sequence
from blinker import signal
from .datasets import InMemoryDataset, GeneratorDataset
from .model_wrappers import OracleWrapper
from .samplers import ConditionalStartSampler, VarianceReductionCondition, \
AdaptiveAdditiveSmoothingSampler, AdditiveSmoothingSampler, ModelSampler, \
ConstantVarianceSampler, ConstantTimeSampler, SCSGSampler
from .reweighting import BiasedReweightingPolicy, NoReweightingPolicy
from .utils.functional import ___, compose, partial


class _BaseImportanceTraining(object):
def __init__(self, model, score="gnorm", layer=None):
"""重要性采样训练的基类。

参数：
model: 要训练的 Keras 模型。
"""
# 检查模型是否兼容，并包装模型以支持重要性采样
self._check_model(model)
self.original_model = model
self.model = OracleWrapper(
model,
self.reweighting,
score=score,
layer=layer
)
self.original_model.optimizer = self.model.optimizer
signal("is.sample").connect(self._on_sample)
signal("is.score").connect(self._on_scores)

def _on_sample(self, event):
self._latest_sample_event = event

def _on_scores(self, scores):
self._latest_scores = scores

def _check_model(self, model):
"""检查模型是否使用了 Dropout 和 BatchNorm，并警告可能影响重要性计算。"""
if any(
isinstance(l, (tf.keras.layers.BatchNormalization, Dropout))
for l in model.layers
):
sys.stderr.write(("[NOTICE]: 您正在使用 BatchNormalization 和/或 Dropout。\n"
"这些层可能会影响重要性计算，建议您用 LayerNormalization 或 "
"在测试模式下的 BatchNormalization 以及 L2 正则化替代它们。\n"))

@property
def reweighting(self):
"""控制重要性采样时偏置的重加权策略。"""
raise NotImplementedError()

def sampler(self, dataset, batch_size, steps_per_epoch, epochs):
"""为给定的数据集创建一个新的采样器，以使用重要性采样进行采样。"""
raise NotImplementedError()

def fit(self, x, y, batch_size=32, epochs=1, verbose=1, callbacks=None,
validation_split=0.0, validation_data=None, steps_per_epoch=None,
on_sample=None, on_scores=None):
# 处理验证数据集
if validation_data is not None:
x_train, y_train = x, y
x_test, y_test = validation_data
elif validation_split > 0:
assert validation_split < 1, "验证集比例必须小于 1"
n = int(round(validation_split * len(x)))
idxs = np.arange(len(x))
np.random.shuffle(idxs)
x_train, y_train = x[idxs[n:]], y[idxs[n:]]
x_test, y_test = x[idxs[:n]], y[idxs[:n]]
else:
x_train, y_train = x, y
x_test, y_test = np.empty(shape=(0, x.shape[1])), np.empty(shape=(0, y.shape[1]))

# 创建数据集
dataset = InMemoryDataset(
x_train,
y_train,
x_test,
y_test,
categorical=False # 不需要转换标签
)

return self.fit_dataset(
dataset=dataset,
batch_size=batch_size,
epochs=epochs,
steps_per_epoch=steps_per_epoch,
verbose=verbose,
callbacks=callbacks,
on_sample=on_sample,
on_scores=on_scores
)

def fit_generator(self, train, steps_per_epoch=None, batch_size=32,
epochs=1, verbose=1, callbacks=None,
validation_data=None, validation_steps=None,
on_sample=None, on_scores=None):
"""创建一个 GeneratorDataset 实例，并使用重要性采样训练模型。

参数：
train: 返回 (inputs, targets) 的生成器或可迭代对象。
steps_per_epoch: 每个 epoch 的梯度更新次数。
batch_size: 每个梯度更新的样本数。
epochs: 训练轮数。
verbose: {0, >0}，是否使用 Keras 的进度条回调。
validation_data: 生成器、可迭代对象或 (inputs, targets) 元组。
validation_steps: 如果 validation_data 是生成器，则需要指定此参数。
on_sample: 接受 sampler, idxs, w, scores 的回调函数。
on_scores: 接受 sampler 和 scores 的回调函数。
"""

def infinite_iterator(iterable):
while True:
yield from iterable

# 处理验证数据集
if validation_data is not None:
if isinstance(validation_data, (tuple, list)):
test = validation_data
test_len = None
elif isinstance(validation_data, Iterable):
test = infinite_iterator(validation_data)
try:
test_len = len(validation_data) * batch_size
except TypeError:
test_len = validation_steps * batch_size if validation_steps else None
else:
test = validation_data
test_len = validation_steps * batch_size if validation_steps else None
else:
test = (np.empty(shape=(0, 1)), np.empty(shape=(0, 1)))
test_len = None
logging.info("未提供验证数据")

# 设置训练数据和每个 epoch 的步数
if isinstance(train, Iterable):
if steps_per_epoch is None:
try:
steps_per_epoch = len(train)
except TypeError:
raise ValueError("使用没有 __len__ 方法的生成器时必须设置 steps_per_epoch")
train = infinite_iterator(train)
elif steps_per_epoch is None:
raise ValueError("steps_per_epoch 未设置")

dataset = GeneratorDataset(train, test, test_len)

return self.fit_dataset(
dataset=dataset,
batch_size=batch_size,
epochs=epochs,
steps_per_epoch=steps_per_epoch,
verbose=verbose,
callbacks=callbacks,
on_sample=on_sample,
on_scores=on_scores
)

def fit_dataset(self, dataset, steps_per_epoch=None, batch_size=32,
epochs=1, verbose=1, callbacks=None, on_sample=None,
on_scores=None):
"""在给定的数据集上训练模型。

参数：
dataset: `BaseDataset` 的实例，提供要训练的数据。
steps_per_epoch: int，每个 epoch 的梯度更新次数。
batch_size: 每个梯度更新的样本数。
epochs: 训练轮数。
verbose: {0, >0}，是否使用 Keras 的进度条回调。
callbacks: Keras 回调列表。
on_sample: 接受 sampler, idxs, w, scores 的回调函数。
on_scores: 接受 sampler 和 scores 的回调函数。
"""
try:
if len(dataset.train_data) < batch_size:
raise ValueError("batch_size 大于训练集大小")
except RuntimeError as e:
assert "no size" in str(e)

# 设置 steps_per_epoch
if steps_per_epoch is None:
steps_per_epoch = len(dataset.train_data) // batch_size

# 创建回调列表
self.history = History()
callbacks = [BaseLogger()] + (callbacks or []) + [self.history]
if verbose > 0:
callbacks += [ProgbarLogger(count_mode="steps")]
callbacks = CallbackList(callbacks)
callbacks.set_model(self.original_model)
callbacks.set_params({
"epochs": epochs,
"steps": steps_per_epoch,
"verbose": verbose,
"do_validation": len(dataset.test_data) > 0,
"metrics": self.metrics_names + [
"val_" + n for n in self.metrics_names
]
})

# 创建采样器
sampler = self.sampler(dataset, batch_size, steps_per_epoch, epochs)

# 开始训练循环
epoch = 0
self.original_model.stop_training = False
callbacks.on_train_begin()
while epoch < epochs:
callbacks.on_epoch_begin(epoch)

for step in range(steps_per_epoch):
batch_logs = {"batch": step, "size": batch_size}
callbacks.on_batch_begin(step, batch_logs)

# 在这里进行重要性采样
idxs, (x_batch, y_batch), w = sampler.sample(batch_size)

# 将 x_batch 和 y_batch 转换为 TensorFlow 张量
x_batch = tf.convert_to_tensor(x_batch)
y_batch = tf.convert_to_tensor(y_batch)
w = tf.convert_to_tensor(w, dtype=tf.float32)

# 自定义训练步骤
loss_value, metrics_results, scores = self.train_batch(x_batch, y_batch, w)

# 更新采样器
sampler.update(idxs, scores)

# 记录指标
batch_logs['loss'] = loss_value
for metric_name, metric_value in zip(self.metrics_names[1:], metrics_results):
batch_logs[metric_name] = metric_value
callbacks.on_batch_end(step, batch_logs)

if on_scores is not None and hasattr(self, "_latest_scores"):
on_scores(
sampler,
self._latest_scores
)

if on_sample is not None:
on_sample(
sampler,
self._latest_sample_event["idxs"],
self._latest_sample_event["w"],
self._latest_sample_event["predicted_scores"]
)

if self.original_model.stop_training:
break

# 在每个 epoch 结束时进行评估
epoch_logs = {}
if len(dataset.test_data) > 0:
val_loss, val_metrics = self.evaluate(dataset.test_data)
epoch_logs['val_loss'] = val_loss
for name, value in zip(self.metrics_names[1:], val_metrics):
epoch_logs['val_' + name] = value
callbacks.on_epoch_end(epoch, epoch_logs)
if self.original_model.stop_training:
break
epoch += 1
callbacks.on_train_end()

return self.history

@property
def metrics_names(self):
def name(x):
if isinstance(x, str):
return x
if hasattr(x, "name"):
return x.name
if hasattr(x, "__name__"):
return x.__name__
return str(x)

metrics = self.model.model.metrics or []
return (
["loss"] +
[name(m) for m in metrics]
)

def train_batch(self, x, y, w):
with tf.GradientTape() as tape:
y_pred = self.model.model(x, training=True)

# 获取损失函数
loss_fn = self.model.model.loss

# 处理字符串形式的损失函数
if isinstance(loss_fn, str):
loss_fn = tf.keras.losses.get(loss_fn)
elif isinstance(loss_fn, dict):
loss_fn = tf.keras.losses.get(loss_fn)
elif not callable(loss_fn):
raise ValueError(f"Unsupported loss function type: {type(loss_fn)}")

# 确保损失函数的 reduction 设置为 NONE
if isinstance(loss_fn, tf.keras.losses.Loss):
loss_config = loss_fn.get_config()
per_example_loss_fn = loss_fn.__class__(reduction=tf.keras.losses.Reduction.NONE, **loss_config)
per_sample_losses = per_example_loss_fn(y, y_pred)
else:
per_sample_losses = loss_fn(y, y_pred)

# 使用正确的评分方法
if self.model.score_method == "loss":
scores = per_sample_losses
else:
scores = tf.ones_like(per_sample_losses)

# 计算样本权重和加权损失
weights = self.reweighting.weight_layer()([scores, w])

# **确保两个张量都是 float32**
per_sample_losses = tf.cast(per_sample_losses, tf.float32)
weights = tf.cast(weights, tf.float32)

weighted_losses = per_sample_losses * weights
loss = tf.reduce_mean(weighted_losses)

trainable_vars = self.model.model.trainable_variables
gradients = tape.gradient(loss, trainable_vars)
self.model.model.optimizer.apply_gradients(zip(gradients, trainable_vars))

self.model.model.compiled_metrics.update_state(y, y_pred, sample_weight=weights)
metrics = [metric.result().numpy() for metric in self.model.model.metrics]

return loss.numpy(), metrics, scores.numpy()

def evaluate(self, test_data):
print(f"test_data type: {type(test_data)}, length: {len(test_data)}")
print(f"test_data contents: {test_data}")

x_test, y_test = test_data[:2]
# 只获取前两个元素
x_test, y_test = test_data[:2]

# 如果 x_test 或 y_test 是元组，提取其中的数组
if isinstance(x_test, tuple) or isinstance(x_test, list):
x_test = x_test[0]
if isinstance(y_test, tuple) or isinstance(y_test, list):
y_test = y_test[0]

# 确保 x_test 和 y_test 是 NumPy 数组或张量
x_test = np.asarray(x_test)
y_test = np.asarray(y_test)

y_pred = self.model.model.predict(x_test, batch_size=128)

# 确保 y_test 和 y_pred 是张量
y_test = tf.convert_to_tensor(y_test)
y_pred = tf.convert_to_tensor(y_pred)

loss_value = self.model.model.compiled_loss(y_test, y_pred, regularization_losses=self.model.model.losses)

# 更新指标
self.model.model.compiled_metrics.update_state(y_test, y_pred)
metrics_results = [m.result().numpy() for m in self.model.model.metrics]

return loss_value.numpy(), metrics_results

class _UnbiasedImportanceTraining(_BaseImportanceTraining):
def __init__(self, model, score="gnorm", layer=None):
self._reweighting = BiasedReweightingPolicy(1.0) # 无偏

super(_UnbiasedImportanceTraining, self).__init__(model, score, layer)

@property
def reweighting(self):
return self._reweighting

class ConstantTimeImportanceTraining(_UnbiasedImportanceTraining):
"""使用常时间重要性采样训练模型。

参数：
model: 要训练的 Keras 模型。
score: {"gnorm", "loss", "full_gnorm"}，用于重要性采样的重要性度量。
layer: None 或 int 或 Layer，计算 gnorm 时使用的层。
"""

def __init__(self, model, backward_pass_multiplier=2.0, extra_samples=0.1,
score="gnorm", layer=None):
self._backward_pass_multiplier = backward_pass_multiplier
self._extra_samples = extra_samples

super(ConstantTimeImportanceTraining, self).__init__(
model,
score,
layer
)

def sampler(self, dataset, batch_size, steps_per_epoch, epochs):
return ConstantTimeSampler(
dataset,
self.reweighting,
self.model,
backward_time=self._backward_pass_multiplier,
)
上述是training.py文件
from blinker import signal
from GAS.GAS_for_N_max import GroverOptimizer

def _get_dataset_length(dset, default=1):
"""Return the dataset's training data length and in case the dataset is
uncountable return a defalt value."""
try:
return len(dset.train_data)
except RuntimeError:
return default


import numpy as np
from blinker import signal

def _get_dataset_length(dset, default=1):
"""返回数据集的训练数据长度，如果无法获取，则返回默认值。"""
try:
return len(dset.train_data)
except RuntimeError:
return default


import numpy as np
from blinker import signal

def _get_dataset_length(dset, default=1):
try:
return len(dset.train_data)
except RuntimeError:
return default

class BaseSampler(object):
def __init__(self, dataset, reweighting):
self.dataset = dataset
self.reweighting = reweighting

def _slice_data(self, x, y, idxs):
if isinstance(x, (list, tuple)):
return [xi[idxs] for xi in x], y[idxs]
else:
return x[idxs], y[idxs]

def _send_messages(self, idxs, xy, w, predicted_scores):
signal("is.sample").send({
"idxs": idxs,
"xy": xy,
"w": w,
"predicted_scores": predicted_scores
})

def _get_samples_with_scores(self, batch_size):
raise NotImplementedError()

def sample(self, batch_size):
# Get the importance scores of some samples
idxs1, scores, xy = self._get_samples_with_scores(batch_size)

# Sample from the available ones
if scores is not None:
# Use argsort to sort scores
idxs_sorted = np.argsort(scores)
idxs2 = self._select_samples(idxs_sorted, scores, batch_size)
else:
# If scores is None, randomly select batch_size sample indices
idxs2 = np.random.choice(len(idxs1), batch_size, p=None)

w = self.reweighting.sample_weights(idxs2, scores)

# Make sure we have the data
if xy is None:
xy = self.dataset.train_data[idxs1[idxs2]]
else:
x, y = xy
xy = self._slice_data(x, y, idxs2)

scores = scores[idxs2] if scores is not None else np.ones(batch_size)
self._send_messages(idxs1[idxs2], xy, w, scores)
return idxs1[idxs2], xy, w

def update(self, idxs, results):
pass

def _select_samples(self, idxs_sorted, scores, batch_size):
"""Child classes should implement this method to select sample indices."""
raise NotImplementedError()
import numpy as np



class ConstantTimeSampler(BaseSampler):
def __init__(self, dataset, reweighting, model, batch_size=16,
backward_time=2.0, tau_th=0.9, ratio=0.5, min_a=0.2):
self.condition = VarianceReductionCondition(tau_th)
self.backward_time = backward_time
self.ratio = ratio
self.min_a = min_a
self.model = model
self.batch_size = batch_size
self.N = len(dataset.train_data) if hasattr(dataset, 'train_data') else _get_dataset_length(dataset, default=1)
print(f"Dataset size (self.N): {self.N}")
self.usage_counts = np.zeros(self.N, dtype=int)

super(ConstantTimeSampler, self).__init__(dataset, reweighting)

def sample(self, batch_size):
print(f"Entering sample method")
print(f"batch_size: {batch_size}, self.N: {self.N}")
print(f"Condition satisfied: {self.condition.satisfied}")
print(f"Requested batch_size: {batch_size}")

if batch_size > self.N:
batch_size = min(batch_size, self.N)

scores_selected = None

if self.condition.satisfied:
print("Entering satisfied condition block")
t = (1.0 + self.backward_time) * batch_size
B = 32 # Number of samples to evaluate
idxs = np.random.choice(self.N, B, replace=False)
x, y = self.dataset.train_data[idxs]

# Obtain and process scores
scores = self.model.score(x, y, batch_size=B)
scores = np.nan_to_num(scores, nan=0.0, posinf=1.0, neginf=0.0)
scores = scores + 1e-8

# Normalize scores
if np.max(scores) != np.min(scores):
scores = (scores - np.min(scores)) / (np.max(scores) - np.min(scores))
else:
scores = np.full_like(scores, 0.5)

# Scale to range 0-7
scores = scores * 7 + 1e-8
scores = np.clip(scores, 1e-8, 7)
scores = np.ceil(scores)

print("Normalized and adjusted scores:", scores)

# Calculate standard deviation
std_scores = np.std(scores)
print("Standard deviation of scores:", std_scores)

# Determine divide value
devide = self._determine_devide(std_scores, batch_size)
print(f"Calculated devide value: {devide}")

num_high = devide
num_other = batch_size - devide

# Add checks
if np.any(np.isnan(scores)):
print("Warning: scores contain NaN, falling back to random sampling.")
print("scores:", scores)
return self._random_sample(batch_size)
if np.any(np.isinf(scores)):
print("Warning: scores contain infinity, falling back to random sampling.")
print("scores:", scores)
return self._random_sample(batch_size)

# Use quantum method to select high-score samples
N_Q = int(num_high)
num_value_qubits = 4
num_key_qubits = 5
num_iterations = 3
optimizer = GroverOptimizer(num_value_qubits, num_key_qubits, num_iterations=num_iterations)
result = optimizer.solve(scores, N_Q)
idxs_high = np.array(result.keys)
print(f"Selected indices using quantum method: {idxs_high}")
print(f"Length of idxs_high: {len(idxs_high)}")

if len(idxs_high) == 0:
print("Warning: Failed to get valid sample indices, falling back to random sampling.")
return self._random_sample(batch_size)

# Handle remaining samples
num_other = batch_size - len(idxs_high)
if num_other > 0:
remaining_indices = np.setdiff1d(np.arange(B), idxs_high)
if len(remaining_indices) >= num_other:
idxs_other = np.random.choice(remaining_indices, num_other, replace=False)
else:
idxs_other = np.random.choice(remaining_indices, num_other, replace=True)
else:
idxs_other = np.array([], dtype=int)

# Combine indices
idxs2 = np.concatenate([idxs_high, idxs_other])

# Retrieve selected samples and scores
x, y = self._slice_data(x, y, idxs2)
scores_selected = scores[idxs2]

scores_selected = scores_selected + 1e-8
w = self.reweighting.sample_weights(np.arange(len(scores_selected)), scores_selected)
else:
# Random sampling remains unchanged
idxs = np.random.choice(self.N, batch_size, replace=False)
idxs2 = np.arange(len(idxs))
x, y = self.dataset.train_data[idxs]
scores_selected = np.ones(len(idxs))
w = np.ones((len(idxs), self.reweighting.weight_size))

# Update condition state
self.condition.update(scores_selected)

self._send_messages(idxs[idxs2], (x, y), w, scores_selected)
print(f"Returned batch shapes: x: {x.shape}, y: {y.shape}, w: {w.shape}")
return (idxs[idxs2], (x, y), w)

def _determine_devide(self, std_scores, batch_size):
# Adjust threshold based on standard deviation
if std_scores > 1.0:
devide = int(batch_size * 0.7)
elif std_scores > 0.5:
devide = int(batch_size * 0.5)
else:
devide = int(batch_size * 0.3)
return devide

def _random_sample(self, batch_size):
# Implement random sampling as a fallback strategy
random_indices = np.random.choice(self.N, size=batch_size, replace=False)
x, y = self.dataset.train_data[random_indices]
w = np.ones((batch_size, self.reweighting.weight_size))
return random_indices, (x, y), w
上述是samplers.py文件
#
# Copyright (c) 2017 Idiap Research Institute, http://www.idiap.ch/
# Written by Angelos Katharopoulos <<PRESIDIO_ANONYMIZED_EMAIL_ADDRESS>>
#

from tensorflow.python.keras import backend as K
from tensorflow.python.keras.layers import Layer
import numpy as np


class ReweightingPolicy(object):
"""ReweightingPolicy defines how we weigh the samples given their
importance.

Each policy should provide
1. A layer implementation for use with Keras models
2. A python implementation for use with the samplers
"""
def weight_layer(self):
"""Return a layer that accepts the scores and the return value of the
sample_weights() method and produces new sample weights."""
raise NotImplementedError()

def sample_weights(self, idxs, scores):
raise NotImplementedError()

@property
def weight_size(self):
raise NotImplementedError()
class AdjustedBiasedReweightingPolicy(ReweightingPolicy):
def __init__(self, k=1.0):
self.k = k

def weight_layer(self):
return AdjustedBiasedReweighting(self.k)

def sample_weights(self, idxs, scores):
N = len(scores)
S1 = scores[np.setdiff1d(np.arange(N), idxs)].sum()

return np.tile([float(N), float(S1)], (len(idxs), 1))

@property
def weight_size(self):
return 2


class BiasedReweightingPolicy(ReweightingPolicy):
def __init__(self, k=1.0):
self.k = k

def weight_layer(self):
return ExternalReweighting()

def sample_weights(self, idxs, scores):
N = len(scores)
s = scores # 因为 scores 对应选定的样本
w = s.sum() / N / s
w_hat = w ** self.k
w_hat *= np.dot(w, s) / np.dot(w_hat, s)
return w_hat

@property
def weight_size(self):
return 1

class NoReweightingPolicy(ReweightingPolicy):
"""Set all sample weights to 1."""
def weight_layer(self):
return ExternalReweighting()

def sample_weights(self, idxs, scores):
return np.ones((len(idxs), 1))

@property
def weight_size(self):
return 1


class CorrectingReweightingPolicy(ReweightingPolicy):
"""CorrectingReweightingPolicy aims to scale the sample weights according
to the mistakes of the importance predictor

Arguments
---------
k: float
The bias power used in all other reweighting schemes
"""
def __init__(self, k=1.0):
self.k = k
self._biased_reweighting = BiasedReweightingPolicy(k)

def weight_layer(self):
return CorrectingReweighting()

def sample_weights(self, idxs, scores):
w = self._biased_reweighting.sample_weights(idxs, scores)

return np.hstack([w, scores[idxs][:, np.newaxis]])

@property
def weight_size(self):
return 2


class AdjustedBiasedReweighting(Layer):
def __init__(self, k=1.0, **kwargs):
self.k = k

super(AdjustedBiasedReweighting, self).__init__(**kwargs)

def build(self, input_shape):
assert isinstance(input_shape, list)
assert len(input_shape) == 2
assert input_shape[0][1] == 1
assert input_shape[1][1] == 2

super(AdjustedBiasedReweighting, self).build(input_shape)

def compute_output_shape(self, input_shape):
return input_shape[0]

def call(self, x):
s, s_hat = x

# Compute the variables defined in the class comment
S2 = K.sum(s)
S1 = s_hat[0, 1]
N = s_hat[0, 0]

# Compute the unbiased weights
a2 = (S1 + S2) / N / s

# Compute the biased weights and the scaling factor t
a1 = K.pow(a2, self.k)
sT = K.transpose(s)
t = K.dot(sT, a2) / K.dot(sT, a1)

return K.stop_gradient([a1 * t])[0]


class ExternalReweighting(Layer):
"""Use the provided input as sample weights"""
def build(self, input_shape):
super(ExternalReweighting, self).build(input_shape)

def compute_output_shape(self, input_shape):
return input_shape[0]

def call(self, x):
return K.stop_gradient(x[1])


class CorrectingReweighting(Layer):
def __init__(self, min_decrease=0, max_increase=2, **kwargs):
self.min_decrease = min_decrease
self.max_increase = max_increase

super(CorrectingReweighting, self).__init__(**kwargs)

def build(self, input_shape):
super(CorrectingReweighting, self).build(input_shape)

def compute_output_shape(self, input_shape):
return input_shape[0]

def call(self, x):
s, x1 = x
a = x1[:, :1]
s_hat = x1[:, 1:2]

# Rescale the weights, making sure we mostly scale down
a_hat = a * K.clip(s_hat / s, self.min_decrease, self.max_increase)

# Scale again so that the reported loss is comparable to the other ones
t = 1
#sT = K.transpose(s)
#t = K.dot(sT, a) / K.dot(sT, a_hat)

return K.stop_gradient([a_hat * t])[0]


UNWEIGHTED = NoReweightingPolicy()
UNBIASED = BiasedReweightingPolicy(k=1.0)
上述是reweighting.py文件
#
# Copyright (c) 2017 Idiap Research Institute, http://www.idiap.ch/
# Written by Angelos Katharopoulos <<PRESIDIO_ANONYMIZED_EMAIL_ADDRESS>>
#

from functools import reduce

from tensorflow.python.keras import backend as K
from tensorflow.python.keras import losses
from tensorflow.python.keras.layers import Layer


def _per_sample_loss(loss_function, mask, x):
"""Compute the per sample loss supporting masking and returning always
tensor of shape (batch_size, 1)

Arguments
---------
loss_function: callable
mask: boolean tensor or None
x: list/tuple of inputs to the loss
"""
# Compute the loss
loss = loss_function(*x)

# Apply masking if needed
if mask is not None:
if not isinstance(mask, list):
mask = [mask]
mask = [K.cast(m, K.floatx()) for m in mask if m is not None]
mask = reduce(lambda a, b: a*b, mask)
mask_dims = len(K.int_shape(mask))
mask_mean = K.mean(mask, axis=list(range(1, mask_dims)), keepdims=True)
loss *= mask
loss /= mask_mean

# If the loss has more than 1 dimensions then aggregate the last dimension
dims = len(K.int_shape(loss))
if dims > 1:
loss = K.mean(loss, axis=list(range(1, dims)))

return K.expand_dims(loss)


class LossLayer(Layer):
"""LossLayer outputs the loss per sample

# Arguments
loss: The loss function to use to combine the model output and the
target
"""
def __init__(self, loss, **kwargs):
self.supports_masking = True
self.loss = losses.get(loss)

super(LossLayer, self).__init__(**kwargs)

def compute_mask(self, inputs, input_mask):
return None

def build(self, input_shape):
pass # Nothing to do

super(LossLayer, self).build(input_shape)

def compute_output_shape(self, input_shape):
# We need two inputs X and y
assert len(input_shape) == 2

# (None, 1) because all losses should be scalar
return (input_shape[0][0], 1)

def call(self, x, mask=None):
return _per_sample_loss(self.loss, mask, x)


class GradientNormLayer(Layer):
"""GradientNormLayer aims to output the gradient norm given a list of
parameters (whose gradient to compute) and a loss function to combine the
two inputs.

# Arguments
parameter_list: A list of Keras variables to compute the gradient
norm for
loss: The loss function to use to combine the model output and the
target into a scalar and then compute the gradient norm
fast: If set to True it means we know that the gradient with respect to
each sample only affects one part of the parameter list so we can
use the batch mode to compute the gradient
"""
def __init__(self, parameter_list, loss, fast=False, **kwargs):
self.supports_masking = True
self.parameter_list = parameter_list
self.loss = losses.get(loss)
self.fast = fast

super(GradientNormLayer, self).__init__(**kwargs)

def compute_mask(self, inputs, input_mask):
return None

def build(self, input_shape):
pass # Nothing to do

super(GradientNormLayer, self).build(input_shape)

def compute_output_shape(self, input_shape):
# We get two inputs
assert len(input_shape) == 2

return (input_shape[0][0], 1)

def call(self, x, mask=None):
# x should be an output and a target
assert len(x) == 2

losses = _per_sample_loss(self.loss, mask, x)
if self.fast:
grads = K.sqrt(K.sum([
self._sum_per_sample(K.square(g))
for g in K.gradients(losses, self.parameter_list)
]))
else:
nb_samples = K.shape(losses)[0]
grads = K.map_fn(
lambda i: self._grad_norm(losses[i]),
K.arange(0, nb_samples),
dtype=K.floatx()
)

return K.reshape(grads, (-1, 1))

def _sum_per_sample(self, x):
"""Sum across all the dimensions except the batch dim"""
# Instead we might be able to use x.ndims but there have been problems
# with ndims and Keras so I think len(int_shape()) is more reliable
dims = len(K.int_shape(x))
return K.sum(x, axis=list(range(1, dims)))

def _grad_norm(self, loss):
grads = K.gradients(loss, self.parameter_list)
return K.sqrt(
sum([
K.sum(K.square(g))
for g in grads
])
)

上述是scores.py文件
from tensorflow.python.keras import backend as K
from tensorflow.python.keras.layers import Layer
from tensorflow.python.keras import initializers

from ..utils.tf_custom import tf


class _BaseNormalization(Layer):
"""Implement utility functions for the normalization layers."""
def _moments(self, x, axes):
return (
K.mean(x, axis=axes, keepdims=True),
K.var(x, axis=axes, keepdims=True)
)


class BatchRenormalization(_BaseNormalization):
"""Batch renormalization layer (Sergey Ioffe, 2017).

# Arguments
momentum: Momentum for the moving average
epsilon: Added to variance to avoid divide by 0
rmax: Maximum correction for the variance
dmax: Maximum correction for the bias
"""
def __init__(self, momentum=0.99, epsilon=1e-3, rmax_0=1., rmax_inf=3.,
dmax_0=0., dmax_inf=5., rmax_duration=40000,
dmax_duration=25000, **kwargs):
super(BatchRenormalization, self).__init__(**kwargs)

self.momentum = momentum
self.epsilon = epsilon
self.rmax_0 = rmax_0
self.rmax_inf = rmax_inf
self.rmax_dur = rmax_duration
self.dmax_0 = dmax_0
self.dmax_inf = dmax_inf
self.dmax_dur = dmax_duration

def build(self, input_shape):
dim = input_shape[-1]
if dim is None:
raise ValueError(("The normalization axis should have a "
"defined dimension"))
self.dim = dim

# Trainable part
self.gamma = self.add_weight(
shape=(dim,),
name="gamma",
initializer=initializers.get("ones")
)
self.beta = self.add_weight(
shape=(dim,),
name="beta",
initializer=initializers.get("zeros")
)

# Statistics
self.moving_mean = self.add_weight(
shape=(dim,),
name="moving_mean",
initializer=initializers.get("zeros"),
trainable=False
)
self.moving_sigma = self.add_weight(
shape=(dim,),
name="moving_sigma",
initializer=initializers.get("ones"),
trainable=False
)

# rmax, dmax and steps
self.steps = self.add_weight(
shape=tuple(),
name="steps",
initializer=initializers.get("zeros"),
trainable=False
)
self.rmax = self.add_weight(
shape=tuple(),
name="rmax",
initializer=initializers.Constant(self.rmax_0),
trainable=False
)
self.dmax = self.add_weight(
shape=tuple(),
name="dmax",
initializer=initializers.Constant(self.dmax_0),
trainable=False
)

self.built = True

def _moments(self, x):
axes = range(len(K.int_shape(x))-1)
if K.backend() == "tensorflow":
return tf.nn.moments(x, axes)
else:
# TODO: Maybe the following can be optimized a bit?
mean = K.mean(K.reshape(x, (-1, self.dim)), axis=0)
var = K.var(K.reshape(x, (-1, self.dim)), axis=0)

return mean, var

def _clip(self, x, x_min, x_max):
if K.backend() == "tensorflow":
return tf.clip_by_value(x, x_min, x_max)
else:
return K.maximum(K.minimum(x, x_max), x_min)

def call(self, inputs, training=None):
x = inputs
assert not isinstance(x, list)

# Compute the minibatch statistics
mean, var = self._moments(x)
sigma = K.sqrt(var + self.epsilon)

# If in training phase set rmax, dmax large so that we use the moving
# averages to do the normalization
rmax = K.in_train_phase(self.rmax, K.constant(1e5), training)
dmax = K.in_train_phase(self.dmax, K.constant(1e5), training)

# Compute the corrections based on rmax, dmax
r = K.stop_gradient(self._clip(
sigma/self.moving_sigma,
1./rmax,
rmax
))
d = K.stop_gradient(self._clip(
(mean - self.moving_mean)/self.moving_sigma,
-dmax,
dmax
))

# Actually do the normalization and the rescaling
xnorm = ((x-mean)/sigma)*r + d
y = self.gamma * xnorm + self.beta

# Add the moving average updates
self.add_update([
K.moving_average_update(self.moving_mean, mean, self.momentum),
K.moving_average_update(self.moving_sigma, sigma, self.momentum)
], x)

# Add the r, d updates
rmax_prog = K.minimum(1., self.steps/self.rmax_dur)
dmax_prog = K.minimum(1., self.steps/self.dmax_dur)
self.add_update([
K.update_add(self.steps, 1),
K.update(
self.rmax,
self.rmax_0 + rmax_prog*(self.rmax_inf-self.rmax_0)
),
K.update(
self.dmax,
self.dmax_0 + dmax_prog*(self.dmax_inf-self.dmax_0)
)
])

# Fix the output's uses learning phase
y._uses_learning_phase = rmax._uses_learning_phase

return y


class LayerNormalization(_BaseNormalization):
"""LayerNormalization is a determenistic normalization layer to replace
BN's stochasticity.

# Arguments
axes: list of axes that won't be aggregated over
"""
def __init__(self, axes=None, bias_axes=[-1], epsilon=1e-3, **kwargs):
super(LayerNormalization, self).__init__(**kwargs)
self.axes = axes
self.bias_axes = bias_axes
self.epsilon = epsilon

def build(self, input_shape):
# Get the number of dimensions and the axes that won't be aggregated
# over
ndims = len(input_shape)
axes = self.axes or []
bias_axes = self.bias_axes or [-1]

# Figure out the shape of the statistics
gamma_shape = [1]*ndims
beta_shape = [1]*ndims
for ax in axes:
gamma_shape[ax] = input_shape[ax]
for ax in bias_axes:
beta_shape[ax] = input_shape[ax]

# Figure out the axes we will aggregate over accounting for negative
# axes
self.reduction_axes = [
ax for ax in range(ndims)
if ax > 0 and (ax+ndims) % ndims not in axes
]

# Create trainable variables
self.gamma = self.add_weight(
shape=gamma_shape,
name="gamma",
initializer=initializers.get("ones")
)
self.beta = self.add_weight(
shape=beta_shape,
name="beta",
initializer=initializers.get("zeros")
)

self.built = True

def call(self, inputs):
x = inputs
assert not isinstance(x, list)

# Compute the per sample statistics
mean, var = self._moments(x, self.reduction_axes)
std = K.sqrt(var + self.epsilon)

return self.gamma*(x-mean)/std + self.beta


class StatsBatchNorm(_BaseNormalization):
"""Use the accumulated statistics for batch norm instead of computing them
for each minibatch.

# Arguments
momentum: Momentum for the moving average
epsilon: Added to variance to avoid divide by 0
"""
def __init__(self, momentum=0.99, epsilon=1e-3, update_stats=False,
**kwargs):
super(StatsBatchNorm, self).__init__(**kwargs)

self.momentum = momentum
self.epsilon = epsilon
self.update_stats = update_stats

def build(self, input_shape):
dim = input_shape[-1]
if dim is None:
raise ValueError(("The normalization axis should have a "
"defined dimension"))
self.dim = dim

# Trainable part
self.gamma = self.add_weight(
shape=(dim,),
name="gamma",
initializer=initializers.get("ones")
)
self.beta = self.add_weight(
shape=(dim,),
name="beta",
initializer=initializers.get("zeros")
)

# Statistics
self.moving_mean = self.add_weight(
shape=(dim,),
name="moving_mean",
initializer=initializers.get("zeros"),
trainable=False
)
self.moving_variance = self.add_weight(
shape=(dim,),
name="moving_variance",
initializer=initializers.get("ones"),
trainable=False
)

self.built = True

def call(self, inputs, training=None):
x = inputs
assert not isinstance(x, list)

# Do the normalization and the rescaling
xnorm = K.batch_normalization(
x,
self.moving_mean,
self.moving_variance,
self.beta,
self.gamma,
epsilon=self.epsilon
)

# Compute and update the minibatch statistics
if self.update_stats:
mean, var = self._moments(x, axes=range(len(K.int_shape(x))-1))
self.add_update([
K.moving_average_update(self.moving_mean, mean, self.momentum),
K.moving_average_update(self.moving_variance, var, self.momentum)
], x)

return xnorm


class GroupNormalization(_BaseNormalization):
"""GroupNormalization is an improvement to LayerNormalization presented in
https://arxiv.org/abs/1803.08494.

#Arguments
G: The channels per group in the normalization
epsilon: Added to the variance to avoid division with 0
"""
def __init__(self, G=32, epsilon=1e-3, **kwargs):
super(GroupNormalization, self).__init__(**kwargs)

self.G = G
self.epsilon = epsilon

def build(self, input_shape):
# Get the number of dimensions and the channel axis
ndims = len(input_shape)
channel_axis = -1 if K.image_data_format() == "channels_last" else -3
shape = [1]*ndims
shape[channel_axis] = input_shape[channel_axis]

# Make sure everything is in order
assert None not in shape, "The channel axis must be defined"
assert shape[channel_axis] % self.G == 0, ("The channels must be "
"divisible by the number "
"of groups")

# Create the trainable weights
self.gamma = self.add_weight(
shape=shape,
name="gamma",
initializer=initializers.get("ones")
)
self.beta = self.add_weight(
shape=shape,
name="beta",
initializer=initializers.get("zeros")
)

# Keep the channel axis for later use
self.channel_axis = channel_axis
self.ndims = ndims

# We 're done
self.built = True

def call(self, inputs):
x = inputs
assert not isinstance(x, list)

# Get the shapes
G = self.G
channel_axis = self.channel_axis
ndims = self.ndims
original_shape = K.shape(x)
shape = [
original_shape[i]
for i in range(ndims)
]
if channel_axis == -1:
shape[channel_axis] //= G
shape.append(G)
axes = sorted([
ndims + channel_axis - i
for i in range(3)
])
else:
shape[channel_axis] //= G
shape.insert(channel_axis, G)
axes = sorted([
ndims + channel_axis + i
for i in range(3)
])

# Do the group norm
x = K.reshape(x, shape)
mean, var = self._moments(x, axes)
std = K.sqrt(var + self.epsilon)
x = (x - mean)/std
x = K.reshape(x, original_shape)

return self.gamma*x + self.beta
上述是normalization.py文件
from functools import reduce
from tensorflow.python.keras import backend as K
from tensorflow.python.keras.layers import Layer
from tensorflow.python.keras.metrics import categorical_accuracy, binary_accuracy, \
get as get_metric, sparse_categorical_accuracy
from ..utils.functional import compose
class MetricLayer(Layer):

def __init__(self, metric_func, **kwargs):
self.supports_masking = True
self.metric_func = metric_func

super(MetricLayer, self).__init__(**kwargs)

def compute_mask(self, inputs, input_mask):
return None

def build(self, input_shape):
# Special care for accuracy because keras treats it specially
try:
if "acc" in self.metric_func:
self.metric_func = self._generic_accuracy
except TypeError:
pass # metric_func is not a string
self.metric_func = compose(K.expand_dims, get_metric(self.metric_func))

super(MetricLayer, self).build(input_shape)

def compute_output_shape(self, input_shape):
# We need two inputs y_true, y_pred
assert len(input_shape) == 2

return (input_shape[0][0], 1)

def call(self, inputs, mask=None):
# Compute the metric
metric = self.metric_func(*inputs)
if K.int_shape(metric)[-1] == 1:
metric = K.squeeze(metric, axis=-1)
if len(K.int_shape(metric)) == 0:
metric = K.ones(K.shape(inputs[0])[0]) * metric

# Apply the mask if needed
if mask is not None:
if not isinstance(mask, list):
mask = [mask]
mask = [K.cast(m, K.floatx()) for m in mask if m is not None]
mask = reduce(lambda a, b: a*b, mask)
metric *= mask
metric /= K.mean(mask, axis=-1, keepdims=True)

# Make sure that the tensor returned is (None, 1)
dims = len(K.int_shape(metric))
if dims > 1:
metric = K.mean(metric, axis=list(range(1, dims)))

return K.expand_dims(metric)

@staticmethod
def _generic_accuracy(y_true, y_pred):
if K.int_shape(y_pred)[1] == 1:
return binary_accuracy(y_true, y_pred)
if K.int_shape(y_true)[-1] == 1:
return sparse_categorical_accuracy(y_true, y_pred)

return categorical_accuracy(y_true, y_pred)


class TripletLossLayer(Layer):

def __init__(self, **kwargs):
super(TripletLossLayer, self).__init__(**kwargs)

def build(self, input_shape):
assert not isinstance(input_shape, list)
self.N = input_shape[1] // 3
self.built = True

def compute_output_shape(self, input_shape):
assert not isinstance(input_shape, list)
return (input_shape[0], 1)

def call(self, x):
N = self.N

xa = x[:, :N]
xp = x[:, N:2*N]
xn = x[:, 2*N:]

xa = xa / K.sqrt(K.sum(xa**2, axis=1, keepdims=True))
xp = xp / K.sqrt(K.sum(xp**2, axis=1, keepdims=True))
xn = xn / K.sqrt(K.sum(xn**2, axis=1, keepdims=True))

dn = K.sum(K.square(xa - xn), axis=1, keepdims=True)
dp = K.sum(K.square(xa - xp), axis=1, keepdims=True)

return dn - dp
上述是mertics.py文件
上述为项目1
重要性采样原文献

from __future__ import print_function
import keras
from keras.models import Sequential
from keras.layers import Activation, Dense, Flatten, Dropout
from keras.layers import Conv2D, MaxPooling2D
from keras.regularizers import l2
from keras.preprocessing.image import ImageDataGenerator
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelBinarizer
import numpy as np
import os
import scipy.io
from skimage.io import imread
from skimage.transform import resize

from importance_sampling.training import ConstantTimeImportanceTraining
from example_utils import get_parser


def load_data(data_dir, img_size=(224, 224)):
    # 加载图像标签
    labels = scipy.io.loadmat(os.path.join(data_dir, 'imagelabels.mat'))['labels'][0]
    num_classes = 102
    img_rows, img_cols = img_size

    # 图像文件夹路径
    img_dir = os.path.join(data_dir, 'jpg')

    # 读取所有图像文件名
    img_files = sorted([os.path.join(img_dir, f'image_{i:05d}.jpg') for i in range(1, len(labels) + 1)])

    # 初始化数据和标签数组
    x_data = []
    y_data = []

    for img_file, label in zip(img_files, labels):
        img = imread(img_file)
        img = resize(img, (img_rows, img_cols), anti_aliasing=True)
        x_data.append(img)
        y_data.append(label - 1)  # 标签从 1 开始，需要减去 1 使其从 0 开始

    x_data = np.array(x_data)
    y_data = np.array(y_data)

    # 标准化数据
    x_data = x_data.astype('float32') / 255.0

    # 将类标签转换为 one-hot 编码
    lb = LabelBinarizer()
    y_data = lb.fit_transform(y_data)

    # 将数据分为训练集和测试集
    x_train, x_test, y_train, y_test = train_test_split(x_data, y_data, test_size=0.2, random_state=42)

    return x_train, x_test, y_train, y_test

if __name__ == "__main__":
    parser = get_parser("在Oxford 102花朵数据集上训练CNN")
    args = parser.parse_args()

    batch_size = 128
    num_classes = 102  # Oxford 102花朵数据集中的类别数
    epochs = 50  # 增加训练 epochs 数量

    # 设置数据目录
    data_dir = 'C:\\Users\\liuwenjing\\GITHUB\\importance-sampling-master\\importance-sampling-master\\102flowers'

    x_train, x_test, y_train, y_test = load_data(data_dir)

    input_shape = x_train.shape[1:]

    print('x_train shape:', x_train.shape)
    print(x_train.shape[0], '训练样本')
    print(x_test.shape[0], '测试样本')

    # 使用数据增强
    datagen = ImageDataGenerator(
        rotation_range=20,
        width_shift_range=0.2,
        height_shift_range=0.2,
        horizontal_flip=True
    )
    datagen.fit(x_train)

    model = Sequential()
    model.add(Conv2D(32, kernel_size=(3, 3), activation='relu', kernel_regularizer=l2(1e-5), input_shape=input_shape))
    model.add(Conv2D(64, (3, 3), activation='relu', kernel_regularizer=l2(1e-5)))
    model.add(MaxPooling2D(pool_size=(2, 2)))
    model.add(Dropout(0.25))
    model.add(Conv2D(128, (3, 3), activation='relu', kernel_regularizer=l2(1e-5)))
    model.add(MaxPooling2D(pool_size=(2, 2)))
    model.add(Dropout(0.25))
    model.add(Flatten())
    model.add(Dense(512, activation='relu', kernel_regularizer=l2(1e-5)))
    model.add(Dropout(0.5))
    model.add(Dense(num_classes, kernel_regularizer=l2(1e-5)))
    model.add(Activation('softmax'))

    model.compile(loss=keras.losses.categorical_crossentropy, optimizer=keras.optimizers.Adam(), metrics=['accuracy'])

    if args.importance_training:
        wrapped = ConstantTimeImportanceTraining(model)
    else:
        wrapped = model

    # 使用数据增强生成器进行训练
    wrapped.fit(datagen.flow(x_train, y_train, batch_size=batch_size), epochs=epochs, verbose=1, validation_data=(x_test, y_test))

    score = model.evaluate(x_test, y_test, verbose=0)
    print('测试损失:', score[0])
    print('测试准确率:', score[1])
上述是102flowers.py文件
#
# Copyright (c) 2017 Idiap Research Institute, http://www.idiap.ch/
# Written by Angelos Katharopoulos <<PRESIDIO_ANONYMIZED_EMAIL_ADDRESS>>
#

import sys

from keras.callbacks import BaseLogger, CallbackList, History, ProgbarLogger
from keras.layers import Dropout, BatchNormalization
from keras.utils import Sequence
import numpy as np
from blinker import signal

from .datasets import InMemoryDataset, GeneratorDataset
from .model_wrappers import OracleWrapper, SVRGWrapper
from .samplers import ConditionalStartSampler, VarianceReductionCondition, \
    AdaptiveAdditiveSmoothingSampler, AdditiveSmoothingSampler, ModelSampler, \
    LSTMSampler, ConstantVarianceSampler, ConstantTimeSampler, SCSGSampler
from .reweighting import BiasedReweightingPolicy, NoReweightingPolicy
from .utils.functional import ___, compose, partial


class _BaseImportanceTraining(object):
    def __init__(self, model, score="gnorm", layer=None):
        """Abstract base class for training a model using importance sampling.

        Arguments
        ---------
            model: The Keras model to train
        """
        # Wrap and transform the model so that it outputs the importance scores
        # and can be used in an importance sampling training scheme
        self._check_model(model)
        self.original_model = model
        self.model = OracleWrapper(
            model,
            self.reweighting,
            score=score,
            layer=layer
        )
        self.original_model.optimizer = self.model.optimizer
        signal("is.sample").connect(self._on_sample)
        signal("is.score").connect(self._on_scores)

    def _on_sample(self, event):
        self._latest_sample_event = event

    def _on_scores(self, scores):
        self._latest_scores = scores

    def _check_model(self, model):
        """Check if the model uses Dropout and BatchNorm and warn that it may
        affect the importance calculations."""
        if any(
            isinstance(l, (BatchNormalization, Dropout))
            for l in model.layers
        ):
            sys.stderr.write(("[NOTICE]: You are using BatchNormalization "
                              "and/or Dropout.\nThose layers may affect the "
                              "importance calculations and you are advised "
                              "to exchange them for LayerNormalization or "
                              "BatchNormalization in test mode and L2 "
                              "regularization.\n"))

    @property
    def reweighting(self):
        """The reweighting policy that controls the bias of the estimation when
        using importance sampling."""
        raise NotImplementedError()

    def sampler(self, dataset, batch_size, steps_per_epoch, epochs):
              raise NotImplementedError()

    def fit(self, x, y, batch_size=32, epochs=1, verbose=1, callbacks=None,
            validation_split=0.0, validation_data=None, steps_per_epoch=None,
            on_sample=None, on_scores=None):
               # Create two data tuples from the given x, y, validation_*
        if validation_data is not None:
            x_train, y_train = x, y
            x_test, y_test = validation_data

        elif validation_split > 0:
            assert validation_split < 1, "100% of the data used for testing"
            n = int(round(validation_split * len(x)))
            idxs = np.arange(len(x))
            np.random.shuffle(idxs)
            x_train, y_train = x[idxs[n:]], y[idxs[n:]]
            x_test, y_test = x[idxs[:n]], y[idxs[:n]]

        else:
            x_train, y_train = x, y
            x_test, y_test = np.empty(shape=(0, 1)), np.empty(shape=(0, 1))

        # Make the dataset to train on
        dataset = InMemoryDataset(
            x_train,
            y_train,
            x_test,
            y_test,
            categorical=False  # this means use the targets as is
        )

        return self.fit_dataset(
            dataset=dataset,
            batch_size=batch_size,
            epochs=epochs,
            steps_per_epoch=steps_per_epoch,
            verbose=verbose,
            callbacks=callbacks,
            on_sample=on_sample,
            on_scores=on_scores
        )

    def fit_generator(self, train, steps_per_epoch=None, batch_size=32,
                      epochs=1, verbose=1, callbacks=None,
                      validation_data=None, validation_steps=None,
                      on_sample=None, on_scores=None):
        """Create a GeneratorDataset instance and train the model using
        importance sampling for a given number of epochs.

        NOTICE: This method may not be supported by all importance training
        classes and may result in NotImplementedError()

        Arguments
        ---------
            train: A generator that returns tuples (inputs, targets)
            steps_per_epoch: int, number of gradient updates before considering
                             an epoch has passed
            batch_size: int, the number of samples per gradient update (ideally
                             set to the number of items returned by the
                             generator at each call)
            epochs: int, multiplied by steps_per_epoch denotes the total number
                    of gradient updates
            verbose: {0, >0}, whether to use the progress bar Keras callback
            validation_data: generator or tuple (inputs, targets)
            validation_steps: None or int, used only if validation_data is a
                              generator
            on_sample: callable that accepts the sampler, idxs, w, scores
            on_scores: callable that accepts the sampler and scores
        """
        def sequence_gen(seq):
            i = 0
            while True:
                yield seq[i]
                i = (i+1) % len(seq)

        # Create the validation data to pass to the GeneratorDataset
        if validation_data is not None:
            if isinstance(validation_data, (tuple, list)):
                test = validation_data
                test_len = None
            elif isinstance(validation_data, Sequence):
                test = sequence_gen(validation_data)
                test_len = len(validation_data) * batch_size
            else:
                test = validation_data
                test_len = validation_steps * batch_size
        else:
            test = (np.empty(shape=(0, 1)), np.empty(shape=(0, 1)))
            test_len = None

        if isinstance(train, Sequence):
            if steps_per_epoch is None:
                steps_per_epoch = len(train)
            train = sequence_gen(train)
        elif steps_per_epoch is None:
            raise ValueError("steps_per_epoch is not set")

        dataset = GeneratorDataset(train, test, test_len)

        return self.fit_dataset(
            dataset=dataset,
            batch_size=batch_size,
            epochs=epochs,
            steps_per_epoch=steps_per_epoch,
            verbose=verbose,
            callbacks=callbacks,
            on_sample=on_sample,
            on_scores=on_scores
        )

    def fit_dataset(self, dataset, steps_per_epoch=None, batch_size=32,
                    epochs=1, verbose=1, callbacks=None, on_sample=None,
                    on_scores=None):
        """Train the model on the given dataset for a given number of epochs.

        Arguments
        ---------
            dataset: Instance of `BaseDataset` that provides the data
                     to train on.
            steps_per_epoch: int or None, number of gradient updates before
                             considering an epoch has passed. If None it is set
                             to be `len(dataset.train_data) / batch_size`.
            batch_size: int, number of samples per gradient update
            epochs: int, number of times to iterate `steps_per_epoch` times
            verbose: {0, >0}, whether to employ the progress bar Keras
                     callback or not
            callbacks: list of Keras callbacks to be called during training
            on_sample: callable that accepts the sampler, idxs, w, scores
            on_scores: callable that accepts the sampler and scores
        """
        try:
            if len(dataset.train_data) < batch_size:
                raise ValueError(("The model cannot be trained with "
                                  "batch_size > training set"))
        except RuntimeError as e:
            assert "no size" in str(e)

        # Set steps_per_epoch properly
        if steps_per_epoch is None:
            steps_per_epoch = len(dataset.train_data) // batch_size

        # Create the callbacks list
        self.history = History()
        callbacks = [BaseLogger()] + (callbacks or []) + [self.history]
        if verbose > 0:
            callbacks += [ProgbarLogger(count_mode="steps")]
        callbacks = CallbackList(callbacks)
        callbacks.set_model(self.original_model)
        callbacks.set_params({
            "epochs": epochs,
            "steps": steps_per_epoch,
            "verbose": verbose,
            "do_validation": len(dataset.test_data) > 0,
            "metrics": self.metrics_names + [
                "val_" + n for n in self.metrics_names
            ]
        })

        # Create the sampler
        sampler = self.sampler(dataset, batch_size, steps_per_epoch, epochs)

        # Start the training loop
        epoch = 0
        self.original_model.stop_training = False
        callbacks.on_train_begin()
        while epoch < epochs:
            callbacks.on_epoch_begin(epoch)
            for step in range(steps_per_epoch):
                batch_logs = {"batch": step, "size": batch_size}
                callbacks.on_batch_begin(step, batch_logs)

                # Importance sampling is done here
                idxs, (x, y), w = sampler.sample(batch_size)
                # Train on the sampled data
                loss, metrics, scores = self.model.train_batch(x, y, w)
                # Update the sampler
                sampler.update(idxs, scores)

                values = map(lambda x: x.mean(), [loss] + metrics)
                for l, o in zip(self.metrics_names, values):
                    batch_logs[l] = o
                callbacks.on_batch_end(step, batch_logs)

                if on_scores is not None and hasattr(self, "_latest_scores"):
                    on_scores(
                        sampler,
                        self._latest_scores
                    )

                if on_sample is not None:
                    on_sample(
                        sampler,
                        self._latest_sample_event["idxs"],
                        self._latest_sample_event["w"],
                        self._latest_sample_event["predicted_scores"]
                    )

                if self.original_model.stop_training:
                    break

            # Evaluate now that an epoch passed
            epoch_logs = {}
            if len(dataset.test_data) > 0:
                val = self.model.evaluate(
                    *dataset.test_data[:],
                    batch_size=batch_size
                )
                epoch_logs = {
                    "val_" + l: o
                    for l, o in zip(self.metrics_names, val)
                }
            callbacks.on_epoch_end(epoch, epoch_logs)
            if self.original_model.stop_training:
                break
            epoch += 1
        callbacks.on_train_end()

        return self.history

    @property
    def metrics_names(self):
        def name(x):
            if isinstance(x, str):
                return x
            if hasattr(x, "name"):
                return x.name
            if hasattr(x, "__name__"):
                return x.__name__
            return str(x)
        metrics = self.original_model.metrics or []
        return (
            ["loss"] +
            [name(m) for m in metrics] +
            ["score"]
        )


class _UnbiasedImportanceTraining(_BaseImportanceTraining):
    def __init__(self, model, score="gnorm", layer=None):
        self._reweighting = BiasedReweightingPolicy(1.0)  # no bias

        super(_UnbiasedImportanceTraining, self).__init__(model, score, layer)

    @property
    def reweighting(self):
        return self._reweighting


class ImportanceTraining(_UnbiasedImportanceTraining):
    """Train a model with exact importance sampling.

    Arguments
    ---------
        model: The Keras model to train
        presample: float, the number of samples to presample for scoring
                   given as a factor of the batch size
        tau_th: float or None, the variance reduction threshold after which we
                enable importance sampling
        forward_batch_size: int or None, the batch size to use for the forward
                            pass during scoring
        score: {"gnorm", "loss", "full_gnorm"}, the importance metric to use
               for importance sampling
        layer: None or int or Layer, the layer to compute the gnorm with
    """
    def __init__(self, model, presample=3.0, tau_th=None,
                 forward_batch_size=None, score="gnorm", layer=None):
        self._presample = presample
        self._tau_th = tau_th
        self._forward_batch_size = forward_batch_size

        # Call the parent to wrap the model
        super(ImportanceTraining, self).__init__(model, score, layer)

    def sampler(self, dataset, batch_size, steps_per_epoch, epochs):
        # Configure the sampler
        large_batch = int(self._presample * batch_size)
        forward_batch_size = self._forward_batch_size or batch_size
        # Compute the threshold using eq. 29 in
        # https://arxiv.org/abs/1803.00942
        B = large_batch
        b = batch_size
        tau_th = float(B + 3*b) / (3*b)
        tau_th = self._tau_th if self._tau_th is not None else tau_th

        return ConditionalStartSampler(
            ModelSampler(
                dataset,
                self.reweighting,
                self.model,
                large_batch=large_batch,
                forward_batch_size=forward_batch_size
            ),
            VarianceReductionCondition(tau_th)
        )


class ConstantVarianceImportanceTraining(_UnbiasedImportanceTraining):
    """Train a model faster by keeping the per iteration variance constant but
    decreasing the time.

    Arguments
    ---------
        model: The Keras model to train
        score: {"gnorm", "loss", "full_gnorm"}, the importance metric to use
               for importance sampling
        layer: None or int or Layer, the layer to compute the gnorm with
    """
    def __init__(self, model, backward_time=2.0, extra_samples=0.2,
                 score="gnorm", layer=None):
        self._backward_time = backward_time
        self._extra_samples = extra_samples

        super(ConstantVarianceImportanceTraining, self).__init__(
            model,
            score,
            layer
        )

    def sampler(self, dataset, batch_size, steps_per_epoch, epochs):
        return ConstantVarianceSampler(
            dataset,
            self.reweighting,
            self.model,
            backward_time=self._backward_time,
            extra_samples=self._extra_samples
        )


class ConstantTimeImportanceTraining(_UnbiasedImportanceTraining):
    """Train a model faster by keeping the per iteration time constant but
    improving the quality of the gradients.

    Arguments
    ---------
        model: The Keras model to train
        score: {"gnorm", "loss", "full_gnorm"}, the importance metric to use
               for importance sampling
        layer: None or int or Layer, the layer to compute the gnorm with
    """
    def __init__(self, model, backward_time=2.0, score="gnorm", layer=None):
        self._backward_time = backward_time

        super(ConstantTimeImportanceTraining, self).__init__(
            model,
            score,
            layer
        )

    def sampler(self, dataset, batch_size, steps_per_epoch, epochs):
        return ConstantTimeSampler(
            dataset,
            self.reweighting,
            self.model,
            backward_time=self._backward_time
        )


class BiasedImportanceTraining(_BaseImportanceTraining):
    """Train a model with exact importance sampling using the loss as
    importance.

    Arguments
    ---------
        model: The Keras model to train
        k: float in (-oo, 1], controls the bias of the training that focuses
           the network on the hard examples (see paper)
        smooth: float, influences the sampling distribution towards uniform by
                adding `smooth` to all probabilities and renormalizing
        adaptive_smoothing: bool, If True the `smooth` argument is a percentage
                            of the average training loss
        presample: int, the number of samples to presample for scoring
        forward_batch_size: int, the batch size to use for the forward pass
                            during scoring
    """
    def __init__(self, model, k=0.5, smooth=0.0, adaptive_smoothing=False,
                 presample=256, forward_batch_size=128):
        # Create the reweighting policy
        self._reweighting = BiasedReweightingPolicy(k)

        # Call the parent to wrap the model
        super(BiasedImportanceTraining, self).__init__(model)

        # Create the sampler factory, the workhorse of the whole deal :-)
        adaptive_smoothing_factory = partial(
            AdaptiveAdditiveSmoothingSampler, ___, smooth
        )
        additive_smoothing_factory = partial(
            AdditiveSmoothingSampler, ___, smooth
        )
        self._sampler = partial(
            ModelSampler,
            ___,
            self.reweighting,
            self.model,
            large_batch=presample,
            forward_batch_size=forward_batch_size
        )
        if adaptive_smoothing and smooth > 0:
            self._sampler = compose(adaptive_smoothing_factory, self._sampler)
        elif smooth > 0:
            self._sampler = compose(additive_smoothing_factory, self._sampler)

    @property
    def reweighting(self):
        return self._reweighting

    def sampler(self, dataset, batch_size, steps_per_epoch, epochs):
        return self._sampler(dataset)


class ApproximateImportanceTraining(_BaseImportanceTraining):
    """Train a model with importance sampling using an LSTM with a class
    embedding to predict the importance of the training samples.

    Arguments
    ---------
        model: The Keras model to train
        k: float in (-oo, 1], controls the bias of the training that focuses
           the network on the hard examples (see paper)
        smooth: float, influences the sampling distribution towards uniform by
                adding `smooth` to all probabilities and renormalizing
        adaptive_smoothing: bool, If True the `smooth` argument is a percentage
                            of the average training loss
        presample: int, the number of samples to presample for scoring
    """
    def __init__(self, model, k=0.5, smooth=0.0, adaptive_smoothing=False,
                 presample=2048):
        # Create the reweighting policy
        self._reweighting = BiasedReweightingPolicy(k)

        # Call the parent to wrap the model
        super(ApproximateImportanceTraining, self).__init__(model)

        # Create the sampler factory, the workhorse of the whole deal :-)
        adaptive_smoothing_factory = partial(
            AdaptiveAdditiveSmoothingSampler, ___, smooth
        )
        additive_smoothing_factory = partial(
            AdditiveSmoothingSampler, ___, smooth
        )
        self._sampler = partial(
            LSTMSampler,
            ___,
            self.reweighting,
            presample=presample
        )
        if adaptive_smoothing and smooth > 0:
            self._sampler = compose(adaptive_smoothing_factory, self._sampler)
        elif smooth > 0:
            self._sampler = compose(additive_smoothing_factory, self._sampler)

    @property
    def reweighting(self):
        return self._reweighting

    def sampler(self, dataset, batch_size, steps_per_epoch, epochs):
        return self._sampler(dataset)

    def fit_generator(*args, **kwargs):
        raise NotImplementedError("ApproximateImportanceTraining doesn't "
                                  "support generator training")


class SVRG(_BaseImportanceTraining):
    """Train a model with Stochastic Variance Reduced Gradient descent.

    See [1, 2] for what this trainer implements.

    [1] Johnson, R. and Zhang, T., 2013. Accelerating stochastic gradient
        descent using predictive variance reduction. In Advances in neural
        information processing systems (pp. 315-323).
    [2] Lei, L. and Jordan, M., 2017, April. Less than a Single Pass:
        Stochastically Controlled Stochastic Gradient. In Artificial
        Intelligence and Statistics (pp. 148-156).

    Arguments
    ---------
        model: The Keras model to train
        B: float, the number of samples to use for estimating the average
           gradient (whole dataset used in case of 0). Given as a factor of the
           batch_size.
        B_rate: float, multiply B with B_rate after every update
        B_over_b: int, How many updates to perform before recalculating the
                  average gradient
    """
    def __init__(self, model, B=10., B_rate=1.0, B_over_b=128):
        self._B = B
        self._B_rate = B_rate
        self._B_over_b = B_over_b

        self.original_model = model
        self.model = SVRGWrapper(model)

    @property
    def reweighting(self):
        """SVRG does not need sample weights so it returns the
        NoReweightingPolicy()"""
        return NoReweightingPolicy()

    def sampler(self, dataset, batch_size, steps_per_epoch, epochs):
        """Create the SCSG sampler"""
        B = int(self._B * batch_size)
        
        return SCSGSampler(
            dataset,
            self.reweighting,
            self.model,
            B,
            self._B_over_b,
            self._B_rate
        )
上述是training.py文件
#
# Copyright (c) 2017 Idiap Research Institute, http://www.idiap.ch/
# Written by Angelos Katharopoulos <<PRESIDIO_ANONYMIZED_EMAIL_ADDRESS>>
#

from functools import partial
import sys

from blinker import signal
from keras import backend as K
from keras.layers import Input, Layer, multiply
from keras.models import Model, clone_model
import numpy as np

from .layers import GradientNormLayer, LossLayer, MetricLayer
from .reweighting import UNBIASED
from .utils.functional import compose


def _tolist(x, acceptable_iterables=(list, tuple)):
    if not isinstance(x, acceptable_iterables):
        x = [x]
    return x


def _get_scoring_layer(score, y_true, y_pred, loss="categorical_crossentropy",
                       layer=None, model=None):
    """Get a scoring layer that computes the score for each pair of y_true,
    y_pred"""
    assert score in ["loss", "gnorm", "full_gnorm", "acc"]

    if score == "loss":
        return LossLayer(loss)([
            y_true,
            y_pred
        ])
    elif score == "gnorm":
        return GradientNormLayer(
            layer.output,
            loss,
            fast=True
        )([
            y_true,
            y_pred
        ])
    elif score == "full_gnorm":
        return GradientNormLayer(
            model.trainable_weights,
            loss,
            fast=False
        )([
            y_true,
            y_pred
        ])
    elif score == "acc":
        return LossLayer(categorical_accuracy)([
            y_true,
            y_pred
        ])


class ModelWrapper(object):
    """The goal of the ModelWrapper is to take a NN and add some extra layers
    that produce a score, a loss and the sample weights to perform importance
    sampling."""
    def _iterate_batches(self, x, y, batch_size):
        bs = batch_size
        for s in range(0, len(y), bs):
            yield [xi[s:s+bs] for xi in _tolist(x)], y[s:s+bs]

    def evaluate(self, x, y, batch_size=128):
        result = np.mean(
            np.vstack([
                self.evaluate_batch(xi, yi)
                for xi, yi in self._iterate_batches(x, y, batch_size)
            ]),
            axis=0
        )

        signal("is.evaluation").send(result)
        return result

    def score(self, x, y, batch_size=128):
        bs = batch_size
        result = np.hstack([
            self.score_batch(xi, yi).T
            for xi, yi in self._iterate_batches(x, y, batch_size)
        ]).T

        signal("is.score").send(result)
        return result

    def set_lr(self, lr):
        """Set the learning rate of the wrapped models.

        We try to set the learning rate on a member variable model and a member
        variable small. If we do not find a member variable model we raise a
        NotImplementedError
        """
        try:
            K.set_value(
                self.optimizer.lr,
                lr
            )
        except AttributeError:
            try:
                K.set_value(
                    self.model.optimizer.lr,
                    lr
                )
            except AttributeError:
                raise NotImplementedError()

        try:
            K.set_value(
                self.small.optimizer.lr,
                lr
            )
        except AttributeError:
            pass

    def evaluate_batch(self, x, y):
        raise NotImplementedError()

    def score_batch(self, x, y):
        raise NotImplementedError()

    def train_batch(self, x, y, w):
        raise NotImplementedError()


class ModelWrapperDecorator(ModelWrapper):
    def __init__(self, model_wrapper, implemented_attributes=set()):
        self.model_wrapper = model_wrapper
        self.implemented_attributes = (
            implemented_attributes | set(["model_wrapper"])
        )

    def __getattribute__(self, name):
        _getattr = object.__getattribute__
        implemented_attributes = _getattr(self, "implemented_attributes")
        if name in implemented_attributes:
            return _getattr(self, name)
        else:
            model_wrapper = _getattr(self, "model_wrapper")
            return getattr(model_wrapper, name)


class OracleWrapper(ModelWrapper):
    AVG_LOSS = 0
    LOSS = 1
    WEIGHTED_LOSS = 2
    SCORE = 3
    METRIC0 = 4

    FUSED_ACTIVATION_WARNING = ("[WARNING]: The last layer has a fused "
                                "activation i.e. Dense(..., "
                                "activation=\"sigmoid\").\nIn order for the "
                                "preactivation to be automatically extracted "
                                "use a separate activation layer (see "
                                "examples).\n")

    def __init__(self, model, reweighting, score="loss", layer=None):
        self.reweighting = reweighting
        self.layer = self._gnorm_layer(model, layer)

        # Augment the model with reweighting, scoring etc
        # Save the new model and the training functions in member variables
        self._augment_model(model, score, reweighting)

    def _gnorm_layer(self, model, layer):
        # If we were given a layer then use it directly
        if isinstance(layer, Layer):
            return layer

        # If we were given a layer index extract the layer
        if isinstance(layer, int):
            return model.layers[layer]

        try:
            # Get the last or the previous to last layer depending on wether
            # the last has trainable weights
            skip_one = not bool(model.layers[-1].trainable_weights)
            last_layer = -2 if skip_one else -1

            # If the last layer has trainable weights that means that we cannot
            # automatically extract the preactivation tensor so we have to warn
            # them because they might be missing out or they might not even
            # have noticed
            if last_layer == -1:
                config = model.layers[-1].get_config()
                if config.get("activation", "linear") != "linear":
                    sys.stderr.write(self.FUSED_ACTIVATION_WARNING)

            return model.layers[last_layer]
        except:
            # In case of an error then probably we are not using the gnorm
            # importance
            return None

    def _augment_model(self, model, score, reweighting):
        # Extract some info from the model
        loss = model.loss
        optimizer = model.optimizer.__class__(**model.optimizer.get_config())
        output_shape = model.get_output_shape_at(0)[1:]
        if isinstance(loss, str) and loss.startswith("sparse"):
            output_shape = output_shape[:-1] + (1,)

        # Make sure that some stuff look ok
        assert not isinstance(loss, list)

        # We need to create two more inputs
        #   1. the targets
        #   2. the predicted scores
        y_true = Input(shape=output_shape)
        pred_score = Input(shape=(reweighting.weight_size,))

        # Create a loss layer and a score layer
        loss_tensor = LossLayer(loss)([y_true, model.get_output_at(0)])
        score_tensor = _get_scoring_layer(
            score,
            y_true,
            model.get_output_at(0),
            loss,
            self.layer,
            model
        )

        # Create the sample weights
        weights = reweighting.weight_layer()([score_tensor, pred_score])

        # Create the output
        weighted_loss = weighted_loss_model = multiply([loss_tensor, weights])
        for l in model.losses:
            weighted_loss += l
        weighted_loss_mean = K.mean(weighted_loss)

        # Create the metric layers
        metrics = model.metrics or []
        metrics = [
            MetricLayer(metric)([y_true, model.get_output_at(0)])
            for metric in metrics
        ]

        # Create a model for plotting and providing access to things such as
        # trainable_weights etc.
        new_model = Model(
            inputs=_tolist(model.get_input_at(0)) + [y_true, pred_score],
            outputs=[weighted_loss_model]
        )

        # Build separate on_batch keras functions for scoring and training
        updates = optimizer.get_updates(
            weighted_loss_mean,
            new_model.trainable_weights
        )
        metrics_updates = []
        if hasattr(model, "metrics_updates"):
            metrics_updates = model.metrics_updates
        learning_phase = []
        if weighted_loss_model._uses_learning_phase:
            learning_phase.append(K.learning_phase())
        inputs = _tolist(model.get_input_at(0)) + [y_true, pred_score] + \
            learning_phase
        outputs = [
            weighted_loss_mean,
            loss_tensor,
            weighted_loss,
            score_tensor
        ] + metrics

        train_on_batch = K.function(
            inputs=inputs,
            outputs=outputs,
            updates=updates + model.updates + metrics_updates
        )
        evaluate_on_batch = K.function(
            inputs=inputs,
            outputs=outputs,
            updates=model.state_updates + metrics_updates
        )

        self.model = new_model
        self.optimizer = optimizer
        self.model.optimizer = optimizer
        self._train_on_batch = train_on_batch
        self._evaluate_on_batch = evaluate_on_batch

    def evaluate_batch(self, x, y):
        if len(y.shape) == 1:
            y = np.expand_dims(y, axis=1)
        dummy_weights = np.ones((y.shape[0], self.reweighting.weight_size))
        inputs = _tolist(x) + [y, dummy_weights] + [0]
        outputs = self._evaluate_on_batch(inputs)

        signal("is.evaluate_batch").send(outputs)

        return np.hstack([outputs[self.LOSS]] + outputs[self.METRIC0:])

    def score_batch(self, x, y):
        if len(y.shape) == 1:
            y = np.expand_dims(y, axis=1)
        dummy_weights = np.ones((y.shape[0], self.reweighting.weight_size))
        inputs = _tolist(x) + [y, dummy_weights] + [0]
        outputs = self._evaluate_on_batch(inputs)

        return outputs[self.SCORE].ravel()


    def train_batch(self, x, y, w):
        if len(y.shape) == 1:
            y = np.expand_dims(y, axis=1)

        # train on a single batch
        outputs = self._train_on_batch(_tolist(x) + [y, w, 1])

        # Add the outputs in a tuple to send to whoever is listening
        result = (
            outputs[self.WEIGHTED_LOSS],
            outputs[self.METRIC0:],
            outputs[self.SCORE]
        )
        signal("is.training").send(result)

        return result


class SVRGWrapper(ModelWrapper):
    """Train using SVRG."""
    def __init__(self, model):
        self._augment(model)

    def _augment(self, model):
        # TODO: There is a lot of overlap with the OracleWrapper, merge some
        #       functionality into a separate function or a parent class

        # Extract info from the model
        loss_function = model.loss
        output_shape = model.get_output_shape_at(0)[1:]

        # Create two identical models one with the current weights and one with
        # the snapshot of the weights
        self.model = model
        self._snapshot = clone_model(model)

        # Create the target variable and compute the losses and the metrics
        inputs = [
            Input(shape=K.int_shape(x)[1:])
            for x in _tolist(model.get_input_at(0))
        ]
        model_output = self.model(inputs)
        snapshot_output = self._snapshot(inputs)
        y_true = Input(shape=output_shape)
        loss = LossLayer(loss_function)([y_true, model_output])
        loss_snapshot = LossLayer(loss_function)([y_true, snapshot_output])
        metrics = self.model.metrics or []
        metrics = [
            MetricLayer(metric)([y_true, model_output])
            for metric in metrics
        ]

        # Make a set of variables that will be holding the batch gradient of
        # the snapshot
        self._batch_grad = [
            K.zeros(K.int_shape(p))
            for p in self.model.trainable_weights
        ]

        # Create an optimizer that computes the variance reduced gradients and
        # get the updates
        loss_mean = K.mean(loss)
        loss_snapshot_mean = K.mean(loss_snapshot)
        optimizer, updates = self._get_updates(
            loss_mean,
            loss_snapshot_mean,
            self._batch_grad
        )

        # Create the training function and gradient computation function
        metrics_updates = []
        if hasattr(self.model, "metrics_updates"):
            metrics_updates = self.model.metrics_updates
        learning_phase = []
        if loss._uses_learning_phase:
            learning_phase.append(K.learning_phase())
        inputs = inputs + [y_true] + learning_phase
        outputs = [loss_mean, loss] + metrics

        train_on_batch = K.function(
            inputs=inputs,
            outputs=outputs,
            updates=updates + self.model.updates + metrics_updates
        )
        evaluate_on_batch = K.function(
            inputs=inputs,
            outputs=outputs,
            updates=self.model.state_updates + metrics_updates
        )
        get_grad = K.function(
            inputs=inputs,
            outputs=K.gradients(loss_mean, self.model.trainable_weights),
            updates=self.model.updates
        )

        self.optimizer = optimizer
        self._train_on_batch = train_on_batch
        self._evaluate_on_batch = evaluate_on_batch
        self._get_grad = get_grad

    def _get_updates(self, loss, loss_snapshot, batch_grad):
        model = self.model
        snapshot = self._snapshot
        class Optimizer(self.model.optimizer.__class__):
            def get_gradients(self, *args):
                grad = K.gradients(loss, model.trainable_weights)
                grad_snapshot = K.gradients(
                    loss_snapshot,
                    snapshot.trainable_weights
                )

                return [
                    g - gs + bg
                    for g, gs, bg in zip(grad, grad_snapshot, batch_grad)
                ]

        optimizer = Optimizer(**self.model.optimizer.get_config())
        return optimizer, \
            optimizer.get_updates(loss, self.model.trainable_weights)

    def evaluate_batch(self, x, y):
        outputs = self._evaluate_on_batch(_tolist(x) + [y, 0])
        signal("is.evaluate_batch").send(outputs)

        return np.hstack(outputs[1:])

    def score_batch(self, x, y):
        raise NotImplementedError()

    def train_batch(self, x, y, w):
        outputs = self._train_on_batch(_tolist(x) + [y, 1])

        result = (
            outputs[0],   # mean loss
            outputs[2:],  # metrics
            outputs[1]    # loss per sample
        )
        signal("is.training").send(result)

        return result

    def update_grad(self, sample_generator):
        sample_generator = iter(sample_generator)
        x, y = next(sample_generator)
        N = len(y)
        gradient_sum = self._get_grad(_tolist(x) + [y, 1])
        for g_sum in gradient_sum:
            g_sum *= N
        for x, y in sample_generator:
            grads = self._get_grad(_tolist(x) + [y, 1])
            n = len(y)
            for g_sum, g in zip(gradient_sum, grads):
                g_sum += g*n
            N += len(y)
        for g_sum in gradient_sum:
            g_sum /= N

        K.batch_set_value(zip(self._batch_grad, gradient_sum))
        self._snapshot.set_weights(self.model.get_weights())


class KatyushaWrapper(SVRGWrapper):
    """Implement Katyusha training on top of plain SVRG."""
    def __init__(self, model, t1=0.5, t2=0.5):
        self.t1 = K.variable(t1, name="tau1")
        self.t2 = K.variable(t2, name="tau2")

        super(KatyushaWrapper, self).__init__(model)

    def _get_updates(self, loss, loss_snapshot, batch_grad):
        optimizer = self.model.optimizer
        t1, t2 = self.t1, self.t2
        lr = optimizer.lr

        # create copies and local copies of the parameters
        shapes = [K.int_shape(p) for p in self.model.trainable_weights]
        x_tilde = [p for p in self._snapshot.trainable_weights]
        z = [K.variable(p) for p in self.model.trainable_weights]
        y = [K.variable(p) for p in self.model.trainable_weights]

        # Get the gradients
        grad = K.gradients(loss, self.model.trainable_weights)
        grad_snapshot = K.gradients(
            loss_snapshot,
            self._snapshot.trainable_weights
        )

        # Collect the updates
        p_plus = [
            t1*zi + t2*x_tildei + (1-t1-t2)*yi
            for zi, x_tildei, yi in
            zip(z, x_tilde, y)
        ]
        vr_grad = [
            gi + bg - gsi
            for gi, bg, gsi in zip(grad, grad_snapshot, batch_grad)
        ]
        updates = [
            K.update(yi, xi - lr * gi)
            for yi, xi, gi in zip(y, p_plus, vr_grad)
        ] + [
            K.update(zi,  zi - lr * gi / t1)
            for zi, xi, gi in zip(z, p_plus, vr_grad)
        ] + [
            K.update(p, xi)
            for p, xi in zip(self.model.trainable_weights, p_plus)
        ]

        return optimizer, updates
上述是model_wrappers.py文件
#
# Copyright (c) 2017 Idiap Research Institute, http://www.idiap.ch/
# Written by Angelos Katharopoulos <<PRESIDIO_ANONYMIZED_EMAIL_ADDRESS>>
#

from functools import partial

from blinker import signal
from keras.layers import Dense, Embedding, Flatten, Input, LSTM, Masking, \
    concatenate
from keras.models import Model
import numpy as np


def _get_dataset_length(dset, default=1):
    """Return the dataset's training data length and in case the dataset is
    uncountable return a defalt value."""
    try:
        return len(dset.train_data)
    except RuntimeError:
        return default


class BaseSampler(object):
    """BaseSampler denotes the interface for all the samplers.

    Samplers should provide the rest of the program with data points to train
    on and corresponding relative weights."""
    def __init__(self, dataset, reweighting):
        self.dataset = dataset
        self.reweighting = reweighting

    def _slice_data(self, x, y, idxs):
        if isinstance(x, (list, tuple)):
            return [xi[idxs] for xi in x], y[idxs]
        else:
            return x[idxs], y[idxs]

    def _send_messages(self, idxs, xy, w, predicted_scores):
        signal("is.sample").send({
            "idxs": idxs,
            "xy": xy,
            "w": w,
            "predicted_scores": predicted_scores
        })

    def _get_samples_with_scores(self, batch_size):
        """Child classes should implement this method.

        Arguments
        ---------
        batch_size: int
                    Return at least that many samples
        
        Return
        ------
        idxs: array
              The indices of some samples in the dataset
        scores: array or None
                The predicted importance scores for the corresponding idxs or
                None for uniform sampling
        xy: tuple or None
            Optionally return the data for the corresponding idxs
        """
        raise NotImplementedError()


    def sample(self, batch_size):
        # Get the importance scores of some samples
        idxs1, scores, xy = self._get_samples_with_scores(batch_size)

        # Sample from the available ones
        p = scores / scores.sum() if scores is not None else None
        idxs2 = np.random.choice(len(idxs1), batch_size, p=p)
        w = self.reweighting.sample_weights(idxs2, scores)

        # Make sure we have the data
        if xy is None:
            xy = self.dataset.train_data[idxs1[idxs2]]
        else:
            x, y = xy
            xy = self._slice_data(x, y, idxs2)

        scores = scores[idxs2] if scores is not None else np.ones(batch_size)
        self._send_messages(idxs1[idxs2], xy, w, scores)
        return idxs1[idxs2], xy, w

    def update(self, idxs, results):
        pass


class UniformSampler(BaseSampler):
    """UniformSampler is the simplest possible sampler which samples the
    dataset uniformly."""
    def __init__(self, dataset, reweighting):
        super(UniformSampler, self).__init__(dataset, reweighting)
        # Basically if we don't know the length the indices don't matter so
        # sample batch_size 0s.
        self.idxs = np.arange(_get_dataset_length(self.dataset, default=1))

    def _get_samples_with_scores(self, batch_size):
        return (
            self.idxs,
            None,
            None
        )


class ModelSampler(BaseSampler):
    """ModelSampler uses a model to score the samples and then performs
    importance sampling based on those scores.

    It can be used to implement several training pipelines where the scoring
    model is separately trained or is sampled from the main model or is the
    main model."""
    def __init__(self, dataset, reweighting, model, large_batch=1024,
                 forward_batch_size=128):
        self.model = model
        self.large_batch = large_batch
        self.forward_batch_size = forward_batch_size
        self.N = _get_dataset_length(dataset, default=1)

        super(ModelSampler, self).__init__(dataset, reweighting)

    def _get_samples_with_scores(self, batch_size):
        assert batch_size < self.large_batch

        # Sample a large number of points in random and score them
        idxs = np.random.choice(self.N, self.large_batch)
        x, y = self.dataset.train_data[idxs]
        scores = self.model.score(x, y, batch_size=self.forward_batch_size)

        return (
            idxs,
            scores,
            (x, y)
        )


class CacheSampler(BaseSampler):
    """CacheSampler uses the recent outputs of the model to determine the
    importance of the samples"""
    def __init__(self, dataset, reweighting, staleness=3, cache_prob=0.5,
                 smooth=0.2):
        # Necessary state for implementing the cache sampler
        self._N = N = len(dataset.train_data)
        self._idxs = np.arange(N)
        self._scores = np.ones((N,))
        self._cache = {}
        self._score_sum = 0.0

        # Configuration
        self._staleness = staleness
        self._alpha = cache_prob / (1-cache_prob)
        self._smooth = smooth

        super(CacheSampler, self).__init__(dataset, reweighting)

    def _get_samples_with_scores(self, batch_size):
        return (
            self._idxs,
            self._scores,
            None
        )

    def update(self, idxs, x):
        # Add the new scores to the cache
        for idx, xi in zip(idxs, x):
            self._cache[idx] = (self._smooth + xi, 0)

        # Remove the stale values
        self._score_sum = 0.0
        keys = self._cache.keys()
        for k in keys:
            item = self._cache[k]
            if item[1] > self._staleness:
                self._scores[k] = 1
                del self._cache[k]
            else:
                self._cache[k] = (item[0], item[1]+1)
                self._score_sum += item[0]

        # Recompute the scores for sampling
        N = self._N
        S = self._score_sum
        a = self._alpha
        for k, (s, _) in self._cache.items():
            self._scores[k] = 1 + a * N * s / S


class LSTMSampler(BaseSampler):
    """Use an LSTM to predict the loss based on the previous losses of each
    sample

    Arguments
    ---------

    dataset: The dataset we want to sample from
    presample: int
               Presample that many data points using uniform sampling to
               decrease the complexity
    history: int
             How many scores per data point to keep in history
    smooth: float
            Depending on whether we are using adaptive smoothing or additive we
            either add smooth*mean or simply smooth to each importance before
            sampling
    log: bool
         Do the regression in log space
    adaptive_smooth: bool
                     Smooth adaptively based on the mean of the scores
    forget: float
            A float less than one to used to calculate the mean of the scores
    """
    def __init__(self, dataset, reweighting, presample=2048, history=10, log=False):
        # Initialize the history for every sample
        init = 1.0
        if log:
            init = np.log(init)
        self.history = np.zeros((len(dataset.train_data), history, 1))
        self.history[:, 0, 0] = init
        self.cnts = np.ones(len(dataset.train_data), dtype=np.uint8)

        # Keep some member variables
        self.presample = presample
        self.log = log

        # Create our LSTM model
        x00 = Input(shape=(history, 1))
        x10 = Input(shape=(1,))
        x0 = Masking(mask_value=0.0)(x00)
        x0 = LSTM(32, return_sequences=True, unroll=True)(x0)
        x0 = LSTM(32, unroll=True)(x0)
        x1 = Embedding(dataset.output_size, 32)(x10)
        x1 = Flatten()(x1)
        x = concatenate([x0, x1])
        y = Dense(1)(x)
        self.model = Model(inputs=[x00, x10], outputs=y)
        self.model.compile(optimizer="adam", loss="mse")

        super(LSTMSampler, self).__init__(dataset, reweighting)

    def _to_ids(self, y):
        try:
            if y.shape[1] > 1:
                return np.expand_dims(y.argmax(axis=1), -1)
        except:
            return y

    def _get_samples_with_scores(self, batch_size):
        """Use the LSTM to predict the loss of each sample"""
        # Presample so that we do not run the LSTM for the whole dataset
        idxs = np.random.choice(len(self.history), self.presample)
        x, y = self.dataset.train_data[idxs]

        # Predict normalize and sample
        scores = self.model.predict(
            [self.history[idxs], self._to_ids(y)],
            batch_size=1024
        ).ravel()

        # Perform the regression in logspace if needed
        if self.log:
            np.exp(scores, scores)
        else:
            np.maximum(scores, 1e-6, scores)

        return (
            idxs,
            scores,
            (x, y)
        )

    def update(self, idxs, x):
        # Fetch the classes for the regression
        _, y = self.dataset.train_data[idxs]

        # If we are doing the regression in logspace
        if self.log:
            x = np.log(x)

        # Train the lstm so that it can predict x given the history
        self.model.train_on_batch([self.history[idxs], self._to_ids(y)], x)

        # Update the history to include x
        full = idxs[self.cnts[idxs] == self.history.shape[1]]
        self.history[full] = np.roll(self.history[full], -1, axis=1)
        self.cnts[full] -= 1
        self.history[idxs, self.cnts[idxs], :1] = x
        self.cnts[idxs] += 1


class PerClassGaussian(BaseSampler):
    """Fit a Gaussian per class to predict the losses"""
    def __init__(self, dataset, reweighting, alpha=0.9, presample=2048):
        self.alpha = alpha
        self.presample = presample
        self.mu = np.ones(dataset.output_size) * np.log(dataset.output_size)
        self.variance = np.ones(dataset.output_size)

        super(PerClassGaussian, self).__init__(dataset, reweighting)

    def _get_samples_with_scores(self, batch_size):
        # Presample so that we do not need to compute everything
        # on the whole dataset
        idxs = np.random.choice(len(self.dataset.train_data), self.presample)
        x, y = self.dataset.train_data[idxs]
        yis = y.ravel() if y.shape[1] == 1 else y.argmax(axis=1)

        # Compute the sampling probs for each of the above presampled
        # data points
        scores = self.mu + np.sqrt(np.maximum(self.variance - self.mu**2, 0))
        scores = scores[yis]

        return (
            idxs,
            scores,
            (x, y)
        )

    def update(self, idxs, x):
        # Fetch the classes in order to model per class information
        _, y = self.dataset.train_data[idxs]
        yis = y.ravel() if y.shape[1] == 1 else y.argmax(axis=1)

        # Update the mean and variance one by one
        # TODO: Improve the following implementation
        for xi, yi in zip(x, yis):
            d = (1.0 - self.alpha) * xi
            self.mu[yi] = self.alpha * self.mu[yi] + d
            self.variance[yi] = self.alpha * self.variance[yi] + d * xi


class LSTMComparisonSampler(BaseSampler):
    """Compare LSTM and Model scores on a fixed presampled subset of the
    training data"""
    def __init__(self, dataset, lstm, model, subset=1024):
        self._idxs = np.random.choice(len(dataset.train_data), subset)
        self._x, self._y= dataset.train_data[self._idxs]
        self.lstm = lstm
        self.model = model

    def _get_samples_with_scores(self, batch_size):
        s1 = self.model.model.score(self._x, self._y)
        s2 = self.lstm.model.predict(
            [self.lstm.history[self._idxs], self.lstm._to_ids(self._y)],
            batch_size=1024
        ).ravel()
        signal("is.lstm_comparison_sampler.scores").send(zip(s1, s2))

        return self.lstm._get_samples_with_scores(batch_size)

    def update(self, idxs, x):
        return self.lstm.update(idxs, x)


class SamplerDecorator(BaseSampler):
    """Just decorate another sampler.

    Arguments
    ---------
    sampler: BaseSampler
             The sampler being decorated
    """
    def __init__(self, sampler):
        self.sampler = sampler

        super(SamplerDecorator, self).__init__(
            sampler.dataset,
            sampler.reweighting
        )

    def _get_samples_with_scores(self, batch_size):
        raise NotImplementedError()

    def update(self, idxs, results):
        self.sampler.update(idxs, results)

    @property
    def model(self):
        """Expose the model attribute of the decorated sampler if one exists."""
        return self.sampler.model


class AdditiveSmoothingSampler(SamplerDecorator):
    """Add a constant to all the importance scores in order to smooth them
    towards uniform

    Arguments
    ---------
    sampler: BaseSampler
             The sampler being decorated
    c: float
       A constant to add to every importance score
    """
    def __init__(self, sampler, c=1.0):
        self.c = c
        super(AdditiveSmoothingSampler, self).__init__(sampler)

    def _get_samples_with_scores(self, batch_size):
        idxs, scores, xy = self.sampler._get_samples_with_scores(batch_size)

        return (
            idxs,
            scores + self.c,
            xy
        )


class AdaptiveAdditiveSmoothingSampler(SamplerDecorator):
    """Add a percentage of the moving average of the predicted importance
    scores to smooth them towards uniform.

    Arguments
    ---------
    sampler: BaseSampler
             The sampler being decorated
    percentage: float
                Multiplied by the moving average of the importance scores to
                add to each score to smooth it
    forget: float
            Used to compute the exponential moving average mu = forget * mu +
            (1-forget) * mu_new
    """
    def __init__(self, sampler, percentage=0.5, forget=0.9):
        self.percentage = percentage
        self.forget = forget
        self.mu = 1.0  # it could be 0 it doesn't really matter
        super(AdaptiveAdditiveSmoothingSampler, self).__init__(sampler)

    def _get_samples_with_scores(self, batch_size):
        idxs, scores, xy = self.sampler._get_samples_with_scores(batch_size)

        self.mu = self.forget * self.mu + (1 - self.forget) * scores.mean()

        return (
            idxs,
            scores + self.percentage * self.mu,
            xy
        )


class PowerSmoothingSampler(SamplerDecorator):
    """Raise the importance scores to a power (less than 1) to smooth them
    towards uniform.
    
    Arguments
    ---------
    sampler: BaseSampler
             The sampler being decorated
    power: float
           The power to raise the scores to
    """
    def __init__(self, sampler, power=0.5):
        assert 0 <= power <= 1

        self.power = power
        super(PowerSmoothingSampler, self).__init__(sampler)

    def _get_samples_with_scores(self, batch_size):
        idxs, scores, xy = self.sampler._get_samples_with_scores(batch_size)

        return (
            idxs,
            scores**self.power,
            xy
        )


class ConditionalStartSampler(SamplerDecorator):
    """ConditionalStartSampler samples uniformly untill a condition is met.

    Arguments
    ---------
        sampler: BaseSampler
                 The sampler to be decorated
        condition: Condition
                   Decide if we should start importance sampling or not
    """
    def __init__(self, sampler, condition):
        # create a uniform sampler to sample from when the condition is not met
        self.uniform = UniformSampler(sampler.dataset, sampler.reweighting)
        self.condition = condition

        super(ConditionalStartSampler, self).__init__(sampler)

    def _get_samples_with_scores(self, batch_size):
        if self.condition.satisfied:
            idxs, scores, xy = \
                self.sampler._get_samples_with_scores(batch_size)
            self.condition.update(scores)
        else:
            idxs, scores, xy = \
                self.uniform._get_samples_with_scores(batch_size)
            if scores is None:
                scores = np.ones(len(idxs))

        return (
            idxs,
            scores,
            xy
        )

    def update(self, idxs, scores):
        if not self.condition.previously_satisfied:
            self.condition.update(scores)
        self.sampler.update(idxs, scores)


class Condition(object):
    """An interface for use with the ConditionalStartSampler."""
    @property
    def satisfied(self):
        raise NotImplementedError()

    @property
    def previously_satisfied(self):
        pass  # not necessary

    def update(self, scores):
        pass  # not necessary


class WarmupCondition(Condition):
    """Wait 'warmup' iterations before using importance sampling.

    Arguments
    ---------
        warmup: int
                The number of iterations to wait before starting importance
                sampling
    """
    def __init__(self, warmup=100):
        self._warmup = warmup
        self._iters = 0

    @property
    def satisfied(self):
        return self._iters > self._warmup

    def update(self, scores):
        self._iters += 1


class ExpCondition(Condition):
    """Assume that the scores are created by an exponential distribution and
    sample only if lamda is larger than x.

    Arguments
    ---------
        lambda_th: float
                   When lambda > lambda_th start importance sampling
        momentum: float
                  The momentum to compute the exponential moving average of
                  lambda
    """
    def __init__(self, lambda_th=2.0, momentum=0.9):
        self._lambda_th = lambda_th
        self._lambda = 0.0
        self._previous_lambda = 0.0
        self._momentum = momentum

    @property
    def satisfied(self):
        self._previous_lambda = self._lambda
        return self._lambda > self._lambda_th

    @property
    def previously_satisfied(self):
        return self._previous_lambda > self._lambda_th

    def update(self, scores):
        self._lambda = (
            self._momentum * self._lambda +
            (1-self._momentum) / scores.mean()
        )


class TotalVariationCondition(Condition):
    """Sample from the decorated sampler if the TV of the scores with the
    uniform distribution is larger than a given value.

    Arguments
    ---------
        tv_th: float
               When tv > tv_th start importance sampling
        momentum: float
                  The momentum to compute the exponential moving average of
                  tv
    """
    def __init__(self, tv_th=0.5, momentum=0.9):
        self._tv_th = tv_th
        self._tv = 0.0
        self._previous_tv = 0.0
        self._momentum = momentum

    @property
    def satisfied(self):
        self._previous_tv = self._tv
        return self._tv > self._tv_th

    @property
    def previously_satisfied(self):
        return self._previous_tv > self._tv_th

    def update(self, scores):
        self._previous_tv = self._tv
        new_tv = 0.5 * np.abs(scores/scores.sum() - 1.0/len(scores)).sum()
        self._tv = (
            self._momentum * self._tv +
            (1-self._momentum) * new_tv
        )


class VarianceReductionCondition(Condition):
    """Sample with importance sampling when the variance reduction is larger
    than a threshold. The variance reduction units are in batch size increment.

    Arguments
    ---------
        vr_th: float
               When vr > vr_th start importance sampling
        momentum: float
                  The momentum to compute the exponential moving average of
                  vr
    """
    def __init__(self, vr_th=1.2, momentum=0.9):
        self._vr_th = vr_th
        self._vr = 0.0
        self._previous_vr = 0.0
        self._momentum = momentum

    @property
    def variance_reduction(self):
        return self._vr

    @property
    def satisfied(self):
        self._previous_vr = self._vr
        return self._vr > self._vr_th

    @property
    def previously_satisfied(self):
        return self._previous_vr > self._vr_th

    def update(self, scores):
        u = 1.0/len(scores)
        S = scores.sum()
        if S == 0:
            g = np.array(u)
        else:
            g = scores/S
        new_vr = 1.0 / np.sqrt(1 - ((g-u)**2).sum()/(g**2).sum())
        self._vr = (
            self._momentum * self._vr +
            (1-self._momentum) * new_vr
        )


class ConstantVarianceSampler(BaseSampler):
    """ConstantVarianceSampler uses the VarianceReductionCondition to sample
    less and less points but keep the variance of the gradients constant.
    
    Arguments
    --------
        dataset: The BaseDataset implementation to sample from
        reweighting: Compute the weights to make the sampling unbiased etc.
        model: Used to compute the importance for importance sampling
        backward_time: The slowdown factor of the backward pass in comparison
                       to the forward pass
        extra_samples: Sample that much more than suggested by the predicted
                       variance reduction to account for the approximation
    """
    def __init__(self, dataset, reweighting, model, backward_time=2.0,
                 extra_samples=0.2):
        self.condition = VarianceReductionCondition(
            1.0 / ((backward_time / (1+backward_time)) - extra_samples)
        )
        self.model = model
        self.extra_samples = extra_samples
        self.N = _get_dataset_length(dataset, default=1)

        super(ConstantVarianceSampler, self).__init__(dataset, reweighting)

    def sample(self, batch_size):
        # Sample batch size uniformly at random
        idxs = np.random.choice(self.N, batch_size)
        idxs2 = np.arange(len(idxs))
        x, y = self.dataset.train_data[idxs]
        scores = np.ones(len(idxs))
        w = np.ones((len(idxs), self.reweighting.weight_size))

        # This means that we can get a speedup by backpropagating less
        if self.condition.satisfied:
            f = 1.0 / self.condition.variance_reduction + self.extra_samples
            N = int(f * batch_size)
            scores = self.model.score(x, y, batch_size=batch_size)
            p = scores / scores.sum()
            idxs2 = np.random.choice(len(idxs), N, p=p)
            w = self.reweighting.sample_weights(idxs2, scores)
            x, y = self._slice_data(x, y, idxs2)
            self.condition.update(scores)

        self._send_messages(idxs[idxs2], (x, y), w, scores[idxs2])

        return (
            idxs[idxs2],
            (x, y),
            w
        )

    def update(self, idxs, scores):
        if not self.condition.previously_satisfied:
            self.condition.update(scores)


class ConstantTimeSampler(BaseSampler):
    """ConstantTimeSampler uses the VarianceReductionCondition to increase the
    quality of the gradients while keeping the time per iterations constant.
    
    Arguments
    ---------
    """
    def __init__(self, dataset, reweighting, model, backward_time=2.0,
                 tau_th=2.0, ratio=0.5, min_a=0.2):
        self.condition = VarianceReductionCondition(tau_th)
        self.backward_time = backward_time
        self.ratio = ratio
        self.min_a = min_a
        self.model = model
        self.N = _get_dataset_length(dataset, default=1)

        super(ConstantTimeSampler, self).__init__(dataset, reweighting)

    def sample(self, batch_size):
        # Check whether the condition is satisfied so that we can sample with
        # importance instead of uniformly
        if self.condition.satisfied:
            # compute the available time
            t = (1.0 + self.backward_time) * batch_size
            # t 是每个批次的总时间预算，它包括了正向传播和反向传播的时间。
            # compute the minimum forward-backward batch
            a = max(
                self.min_a * batch_size,
                batch_size / self.condition.variance_reduction
            )
            # a 确保批次大小不小于一个最小比例，并根据方差减少量调整。
            # compute the maximum scored samples
            B = t - (1.0 + self.backward_time)*a
            # split the difference according to ratio keeping the time fixed
            B = int(batch_size + (B - batch_size)*self.ratio)
            # B 是可以评分的最大样本数量，它根据总时间预算和最小批次大小进行调整。
            a = int((t - B) / (1.0 + self.backward_time))

            # do the hippy shake
            idxs = np.random.choice(self.N, B)
            x, y = self.dataset.train_data[idxs]
            scores = self.model.score(x, y, batch_size=batch_size)
            p = scores / scores.sum()
            idxs2 = np.random.choice(B, a, p=p)
            w = self.reweighting.sample_weights(idxs2, scores)
            x, y = self._slice_data(x, y, idxs2)
            self.condition.update(scores)
        else:
            idxs = np.random.choice(self.N, batch_size)
            idxs2 = np.arange(len(idxs))
            x, y = self.dataset.train_data[idxs]
            scores = np.ones(len(idxs))
            w = np.ones((len(idxs), self.reweighting.weight_size))

        self._send_messages(idxs[idxs2], (x, y), w, scores[idxs2])

        return (
            idxs[idxs2],
            (x, y),
            w
        )

    def update(self, idxs, scores):
        if not self.condition.previously_satisfied:
            self.condition.update(scores)


class HistorySampler(ModelSampler):
    """HistorySampler uses the history of the loss to perform importance
    sampling.

    Arguments
    ---------
    dataset: The dataset to sample from
    reweighting: The reweighting scheme
    model: The model to be used for scoring
    recompute: Compute the loss for the whole dataset every recompute batches
    """
    def __init__(self, dataset, reweighting, model, forward_batch_size=128,
                 recompute=2):
        super(HistorySampler, self).__init__(
            dataset,
            reweighting,
            model,
            forward_batch_size=forward_batch_size
        )

        # The configuration of HistorySampler
        self.recompute = recompute

        # Mutable variables holding the state of the sampler
        self._batch = 0
        self._scores = np.ones((len(dataset.train_data),))
        self._unseen = np.ones(len(dataset.train_data), dtype=np.bool)
        self._seen = np.zeros_like(self._unseen)

    def _get_samples_with_scores(self, batch_size):
        return (
            np.arange(len(self._scores)),
            self._scores,
            None
        )

    def update(self, idxs,results):
        # Update the scores of the seen samples
        self._scores[idxs] = results.ravel()
        self._unseen[idxs] = False
        self._seen[idxs] = True
        self._scores[self._unseen] = self._scores[self._seen].mean()

        # Recompute all the scores if needed
        self._batch += 1
        if self._batch % self.recompute == 0:
            for i in range(0, len(self.dataset.train_data), 1024*64):
                x, y = self.dataset.train_data[i:i+1024*64]
                self._scores[i:i+1024*64] = self.model.score(
                    x, y,
                    batch_size=self.forward_batch_size
                ).ravel()
            self._seen[:] = True
            self._unseen[:] = False


class OnlineBatchSelectionSampler(ModelSampler):
    """OnlineBatchSelection is the online batch creation method by Loschchilov
    & Hutter.

    See 'Online Batch Selection for Faster Training of Neural Networks'.

    Arguments
    ---------
    dataset: The dataset to sample from
    reweighting: The reweighting scheme
    model: The model to be used for scoring
    steps_per_epoch: int
                     How many batches to create before considering that an
                     epoch has passed
    recompute: int
               Recompute the scores after r minibatches seen
    s_e: tuple
         Used to compute the sampling probabilities from the ranking
    n_epochs: int
              The number of epochs, used to compute the sampling probabilities
    """
    def __init__(self, dataset, reweighting, model, large_batch=1024,
                 forward_batch_size=128, steps_per_epoch=300, recompute=2,
                 s_e=(1, 1), n_epochs=1):
        super(OnlineBatchSelectionSampler, self).__init__(
            dataset,
            reweighting,
            model,
            large_batch=large_batch,
            forward_batch_size=forward_batch_size
        )

        # The configuration of OnlineBatchSelection
        self.steps_per_epoch = steps_per_epoch
        self.recompute = recompute
        self.s_e = s_e
        self.n_epochs = n_epochs

        # Mutable variables to be updated
        self._batch = 0
        self._epoch = 0
        self._raw_scores = np.ones((len(dataset.train_data),))
        self._scores = np.ones_like(self._raw_scores)
        self._ranks = np.arange(len(dataset.train_data))

    def _get_samples_with_scores(self, batch_size):
        return (
            np.arange(len(self._ranks)),
            self._scores,
            None
        )

    def update(self, idxs, results):
        # Compute the current epoch and the current batch
        self._batch += 1
        self._epoch = 1 + self._batch // self.steps_per_epoch

        # Add the new scores to the raw_scores
        self._raw_scores[idxs] = results.ravel()

        # if it is a new epoch
        if self._batch % self.steps_per_epoch == 0:
            # For the very first batch or every 'recompute' epochs compute the
            # loss across the whole dataset
            if self.recompute > 0 and self._epoch % self.recompute == 0:
                # Extra internal batch size so that we do not load too much
                # data into memory
                scores = []
                for i in range(0, len(self.dataset.train_data), 1024*64):
                    x, y = self.dataset.train_data[i:i+1024*64]
                    scores.append(self.model.score(
                        x, y, batch_size=self.forward_batch_size
                    ))
                self._raw_scores[:] = np.hstack(scores)

            # Sort and recompute the ranks
            N = len(self.dataset.train_data)
            self._ranks[self._raw_scores.argsort()] = np.arange(N)[::-1]

            # Recompute the sample scores from the ranks
            s_e0, s_eend = self.s_e
            n_epochs = self.n_epochs
            s = s_e0 * np.exp(np.log(s_eend/s_e0)/n_epochs) ** self._epoch
            s = 1.0 / np.exp(np.log(s)/N)
            self._scores = s**self._ranks


class SCSGSampler(BaseSampler):
    """Implement [1] using the SVRG model wrapper.

    SCSG is an online version of SVRG especially made for problems with
    difficult to compute batch gradients.

    [1]: Nonconvex Finite-Sum Optimization Via SCSG Methods

    Arguments
    ---------
        dataset: The dataset to sample from
        reweighting: In case we need to reweigh the samples (could be none in
                     this case)
        model: The model wrapper (must be implement the interface of
               SVRGWrapper)
        B: The initial large batch to sample (if None fall back to SVRG)
        B_over_b: The number of minibatches in an iteration
        B_rate: B*B_rate^(epoch) gives B for every epoch
    """
    def __init__(self, dataset, reweighting, model, B=1024, B_over_b=32,
                 B_rate=1.0):
        self.N = _get_dataset_length(dataset, default=1)
        self.B = B or len(dataset.train_data)
        self.B_over_b = B_over_b
        self.B_rate = B_rate
        self._iteration = 0
        self._idxs = np.arange(self.N)
        self._model = model

        super(SCSGSampler, self).__init__(dataset, reweighting)

    def _get_samples_with_scores(self, batch_size):
        if self._iteration % self.B_over_b == 0:
            self._compute_batch_gradient(batch_size)
            self.B *= self.B_rate
            self.B = min(self.B, self.N) if self.N > 1 else self.B
        self._iteration += 1

        return (self._idxs, None, None)

    def _compute_batch_gradient(self, batch_size):
        def batch_gen():
            np.random.shuffle(self._idxs)
            for s in range(0, int(self.B), batch_size):
                yield self.dataset.train_data[self._idxs[s:s+batch_size]]
        self._model.update_grad(batch_gen())
上述是samplers.py文件
#
# Copyright (c) 2017 Idiap Research Institute, http://www.idiap.ch/
# Written by Angelos Katharopoulos <<PRESIDIO_ANONYMIZED_EMAIL_ADDRESS>>
#

from functools import reduce

from keras import backend as K
from keras import objectives
from keras.layers import Layer


def _per_sample_loss(loss_function, mask, x):
    """Compute the per sample loss supporting masking and returning always
    tensor of shape (batch_size, 1)

    Arguments
    ---------
        loss_function: callable
        mask: boolean tensor or None
        x: list/tuple of inputs to the loss
    """
    # Compute the loss
    loss = loss_function(*x)

    # Apply masking if needed
    if mask is not None:
        if not isinstance(mask, list):
            mask = [mask]
        mask = [K.cast(m, K.floatx()) for m in mask if m is not None]
        mask = reduce(lambda a, b: a*b, mask)
        mask_dims = len(K.int_shape(mask))
        mask_mean = K.mean(mask, axis=list(range(1, mask_dims)), keepdims=True)
        loss *= mask
        loss /= mask_mean

    # If the loss has more than 1 dimensions then aggregate the last dimension
    dims = len(K.int_shape(loss))
    if dims > 1:
        loss = K.mean(loss, axis=list(range(1, dims)))

    return K.expand_dims(loss)


class LossLayer(Layer):
    """LossLayer outputs the loss per sample
    
    # Arguments
        loss: The loss function to use to combine the model output and the
              target
    """
    def __init__(self, loss, **kwargs):
        self.supports_masking = True
        self.loss = objectives.get(loss)

        super(LossLayer, self).__init__(**kwargs)

    def compute_mask(self, inputs, input_mask):
        return None

    def build(self, input_shape):
        pass # Nothing to do

        super(LossLayer, self).build(input_shape)

    def compute_output_shape(self, input_shape):
        # We need two inputs X and y
        assert len(input_shape) == 2

        # (None, 1) because all losses should be scalar
        return (input_shape[0][0], 1)

    def call(self, x, mask=None):
        return _per_sample_loss(self.loss, mask, x)


class GradientNormLayer(Layer):
    """GradientNormLayer aims to output the gradient norm given a list of
    parameters (whose gradient to compute) and a loss function to combine the
    two inputs.
    
    # Arguments
        parameter_list: A list of Keras variables to compute the gradient
                        norm for
        loss: The loss function to use to combine the model output and the
              target into a scalar and then compute the gradient norm
        fast: If set to True it means we know that the gradient with respect to
              each sample only affects one part of the parameter list so we can
              use the batch mode to compute the gradient
    """
    def __init__(self, parameter_list, loss, fast=False, **kwargs):
        self.supports_masking = True
        self.parameter_list = parameter_list
        self.loss = objectives.get(loss)
        self.fast = fast

        super(GradientNormLayer, self).__init__(**kwargs)

    def compute_mask(self, inputs, input_mask):
        return None

    def build(self, input_shape):
        pass # Nothing to do

        super(GradientNormLayer, self).build(input_shape)

    def compute_output_shape(self, input_shape):
        # We get two inputs
        assert len(input_shape) == 2

        return (input_shape[0][0], 1)

    def call(self, x, mask=None):
        # x should be an output and a target
        assert len(x) == 2

        losses = _per_sample_loss(self.loss, mask, x)
        if self.fast:
            grads = K.sqrt(sum([
                self._sum_per_sample(K.square(g))
                for g in K.gradients(losses, self.parameter_list)
            ]))
        else:
            nb_samples = K.shape(losses)[0]
            grads = K.map_fn(
                lambda i: self._grad_norm(losses[i]),
                K.arange(0, nb_samples),
                dtype=K.floatx()
            )

        return K.reshape(grads, (-1, 1))

    def _sum_per_sample(self, x):
        """Sum across all the dimensions except the batch dim"""
        # Instead we might be able to use x.ndims but there have been problems
        # with ndims and Keras so I think len(int_shape()) is more reliable
        dims = len(K.int_shape(x))
        return K.sum(x, axis=list(range(1, dims)))

    def _grad_norm(self, loss):
        grads = K.gradients(loss, self.parameter_list)
        return K.sqrt(
            sum([
                K.sum(K.square(g))
                for g in grads
            ])
        )

上述是scores.py文件
#
# Copyright (c) 2017 Idiap Research Institute, http://www.idiap.ch/
# Written by Angelos Katharopoulos <<PRESIDIO_ANONYMIZED_EMAIL_ADDRESS>>
#

from keras import backend as K
from keras.engine import Layer
from keras import initializers

from ..utils.tf import tf


class _BaseNormalization(Layer):
    def _moments(self, x, axes):
        return (
            K.mean(x, axis=axes, keepdims=True),
            K.var(x, axis=axes, keepdims=True)
        )


class BatchRenormalization(_BaseNormalization):
    def __init__(self, momentum=0.99, epsilon=1e-3, rmax_0=1., rmax_inf=3.,
                 dmax_0=0., dmax_inf=5., rmax_duration=40000,
                 dmax_duration=25000, **kwargs):
        super(BatchRenormalization, self).__init__(**kwargs)

        self.momentum = momentum
        self.epsilon = epsilon
        self.rmax_0 = rmax_0
        self.rmax_inf = rmax_inf
        self.rmax_dur = rmax_duration
        self.dmax_0 = dmax_0
        self.dmax_inf = dmax_inf
        self.dmax_dur = dmax_duration

    def build(self, input_shape):
        dim = input_shape[-1]
        if dim is None:
            raise ValueError(("The normalization axis should have a "
                              "defined dimension"))
        self.dim = dim

        # Trainable part
        self.gamma = self.add_weight(
            shape=(dim,),
            name="gamma",
            initializer=initializers.get("ones")
        )
        self.beta = self.add_weight(
            shape=(dim,),
            name="beta",
            initializer=initializers.get("zeros")
        )

        # Statistics
        self.moving_mean = self.add_weight(
            shape=(dim,),
            name="moving_mean",
            initializer=initializers.get("zeros"),
            trainable=False
        )
        self.moving_sigma = self.add_weight(
            shape=(dim,),
            name="moving_sigma",
            initializer=initializers.get("ones"),
            trainable=False
        )

        # rmax, dmax and steps
        self.steps = self.add_weight(
            shape=tuple(),
            name="steps",
            initializer=initializers.get("zeros"),
            trainable=False
        )
        self.rmax = self.add_weight(
            shape=tuple(),
            name="rmax",
            initializer=initializers.Constant(self.rmax_0),
            trainable=False
        )
        self.dmax = self.add_weight(
            shape=tuple(),
            name="dmax",
            initializer=initializers.Constant(self.dmax_0),
            trainable=False
        )

        self.built = True
    def _moments(self, x):
        axes = range(len(K.int_shape(x))-1)
        if K.backend() == "tensorflow":
            return tf.nn.moments(x, axes)
        else:
            # TODO: Maybe the following can be optimized a bit?
            mean = K.mean(K.reshape(x, (-1, self.dim)), axis=0)
            var = K.var(K.reshape(x, (-1, self.dim)), axis=0)
            return mean, var
    def _clip(self, x, x_min, x_max):
        if K.backend() == "tensorflow":
            return tf.clip_by_value(x, x_min, x_max)
        else:
            return K.maximum(K.minimum(x, x_max), x_min)

    def call(self, inputs, training=None):
        x = inputs
        assert not isinstance(x, list)

        # Compute the minibatch statistics
        mean, var = self._moments(x)
        sigma = K.sqrt(var + self.epsilon)

        # If in training phase set rmax, dmax large so that we use the moving
        # averages to do the normalization
        rmax = K.in_train_phase(self.rmax, K.constant(1e5), training)
        dmax = K.in_train_phase(self.dmax, K.constant(1e5), training)

        # Compute the corrections based on rmax, dmax
        r = K.stop_gradient(self._clip(
            sigma/self.moving_sigma,
            1./rmax,
            rmax
        ))
        d = K.stop_gradient(self._clip(
            (mean - self.moving_mean)/self.moving_sigma,
            -dmax,
            dmax
        ))

        # Actually do the normalization and the rescaling
        xnorm = ((x-mean)/sigma)*r + d
        y = self.gamma * xnorm + self.beta

        # Add the moving average updates
        self.add_update([
            K.moving_average_update(self.moving_mean, mean, self.momentum),
            K.moving_average_update(self.moving_sigma, sigma, self.momentum)
        ], x)

        # Add the r, d updates
        rmax_prog = K.minimum(1., self.steps/self.rmax_dur)
        dmax_prog = K.minimum(1., self.steps/self.dmax_dur)
        self.add_update([
            K.update_add(self.steps, 1),
            K.update(
                self.rmax,
                self.rmax_0 + rmax_prog*(self.rmax_inf-self.rmax_0)
            ),
            K.update(
                self.dmax,
                self.dmax_0 + dmax_prog*(self.dmax_inf-self.dmax_0)
            )
        ])

        # Fix the output's uses learning phase
        y._uses_learning_phase = rmax._uses_learning_phase

        return y


class LayerNormalization(_BaseNormalization):
       def __init__(self, axes=None, bias_axes=[-1], epsilon=1e-3, **kwargs):
        super(LayerNormalization, self).__init__(**kwargs)
        self.axes = axes
        self.bias_axes = bias_axes
        self.epsilon = epsilon

    def build(self, input_shape):
        # Get the number of dimensions and the axes that won't be aggregated
        # over
        ndims = len(input_shape)
        axes = self.axes or []
        bias_axes = self.bias_axes or [-1]

        # Figure out the shape of the statistics
        gamma_shape = [1]*ndims
        beta_shape = [1]*ndims
        for ax in axes:
            gamma_shape[ax] = input_shape[ax]
        for ax in bias_axes:
            beta_shape[ax] = input_shape[ax]

        # Figure out the axes we will aggregate over accounting for negative
        # axes
        self.reduction_axes = [
            ax for ax in range(ndims)
            if ax > 0 and (ax+ndims) % ndims not in axes
        ]

        # Create trainable variables
        self.gamma = self.add_weight(
            shape=gamma_shape,
            name="gamma",
            initializer=initializers.get("ones")
        )
        self.beta = self.add_weight(
            shape=beta_shape,
            name="beta",
            initializer=initializers.get("zeros")
        )

        self.built = True

    def call(self, inputs):
        x = inputs
        assert not isinstance(x, list)

        # Compute the per sample statistics
        mean, var = self._moments(x, self.reduction_axes)
        std = K.sqrt(var + self.epsilon)

        return self.gamma*(x-mean)/std + self.beta


class StatsBatchNorm(_BaseNormalization):
    def __init__(self, momentum=0.99, epsilon=1e-3, update_stats=False,
                 **kwargs):
        super(StatsBatchNorm, self).__init__(**kwargs)

        self.momentum = momentum
        self.epsilon = epsilon
        self.update_stats = update_stats

    def build(self, input_shape):
        dim = input_shape[-1]
        if dim is None:
            raise ValueError(("The normalization axis should have a "
                              "defined dimension"))
        self.dim = dim

        # Trainable part
        self.gamma = self.add_weight(
            shape=(dim,),
            name="gamma",
            initializer=initializers.get("ones")
        )
        self.beta = self.add_weight(
            shape=(dim,),
            name="beta",
            initializer=initializers.get("zeros")
        )

        # Statistics
        self.moving_mean = self.add_weight(
            shape=(dim,),
            name="moving_mean",
            initializer=initializers.get("zeros"),
            trainable=False
        )
        self.moving_variance = self.add_weight(
            shape=(dim,),
            name="moving_variance",
            initializer=initializers.get("ones"),
            trainable=False
        )

        self.built = True

    def call(self, inputs, training=None):
        x = inputs
        assert not isinstance(x, list)

        # Do the normalization and the rescaling
        xnorm = K.batch_normalization(
            x,
            self.moving_mean,
            self.moving_variance,
            self.beta,
            self.gamma,
            epsilon=self.epsilon
        )

        # Compute and update the minibatch statistics
        if self.update_stats:
            mean, var = self._moments(x, axes=range(len(K.int_shape(x))-1))
            self.add_update([
                K.moving_average_update(self.moving_mean, mean, self.momentum),
                K.moving_average_update(self.moving_variance, var, self.momentum)
            ], x)

        return xnorm


class GroupNormalization(_BaseNormalization):
        def __init__(self, G=32, epsilon=1e-3, **kwargs):
        super(GroupNormalization, self).__init__(**kwargs)

        self.G = G
        self.epsilon = epsilon

    def build(self, input_shape):
        # Get the number of dimensions and the channel axis
        ndims = len(input_shape)
        channel_axis = -1 if K.image_data_format() == "channels_last" else -3
        shape = [1]*ndims
        shape[channel_axis] = input_shape[channel_axis]

        # Make sure everything is in order
        assert None not in shape, "The channel axis must be defined"
        assert shape[channel_axis] % self.G == 0, ("The channels must be "
                                                   "divisible by the number "
                                                   "of groups")

        # Create the trainable weights
        self.gamma = self.add_weight(
            shape=shape,
            name="gamma",
            initializer=initializers.get("ones")
        )
        self.beta = self.add_weight(
            shape=shape,
            name="beta",
            initializer=initializers.get("zeros")
        )

        # Keep the channel axis for later use
        self.channel_axis = channel_axis
        self.ndims = ndims

        # We 're done
        self.built = True

    def call(self, inputs):
        x = inputs
        assert not isinstance(x, list)

        # Get the shapes
        G = self.G
        channel_axis = self.channel_axis
        ndims = self.ndims
        original_shape = K.shape(x)
        shape = [
            original_shape[i]
            for i in range(ndims)
        ]
        if channel_axis == -1:
            shape[channel_axis] //= G
            shape.append(G)
            axes = sorted([
                ndims + channel_axis - i
                for i in range(3)
            ])
        else:
            shape[channel_axis] //= G
            shape.insert(channel_axis, G)
            axes = sorted([
                ndims + channel_axis + i
                for i in range(3)
            ])

        # Do the group norm
        x = K.reshape(x, shape)
        mean, var = self._moments(x, axes)
        std = K.sqrt(var + self.epsilon)
        x = (x - mean)/std
        x = K.reshape(x, original_shape)

        return self.gamma*x + self.beta
上述是normalization.py文件
#
# Copyright (c) 2017 Idiap Research Institute, http://www.idiap.ch/
# Written by Angelos Katharopoulos <<PRESIDIO_ANONYMIZED_EMAIL_ADDRESS>>
#

from functools import reduce

from keras import backend as K
from keras.engine import Layer
from keras.metrics import categorical_accuracy, binary_accuracy, \
    get as get_metric, sparse_categorical_accuracy

from ..utils.functional import compose


class MetricLayer(Layer):
       def __init__(self, metric_func, **kwargs):
        self.supports_masking = True
        self.metric_func = metric_func

        super(MetricLayer, self).__init__(**kwargs)

    def compute_mask(self, inputs, input_mask):
        return None

    def build(self, input_shape):
        # Special care for accuracy because keras treats it specially
        try:
            if "acc" in self.metric_func:
                self.metric_func = self._generic_accuracy
        except TypeError:
            pass # metric_func is not a string
        self.metric_func = compose(K.expand_dims, get_metric(self.metric_func))

        super(MetricLayer, self).build(input_shape)

    def compute_output_shape(self, input_shape):
        # We need two inputs y_true, y_pred
        assert len(input_shape) == 2

        return (input_shape[0][0], 1)

    def call(self, inputs, mask=None):
        # Compute the metric
        metric = self.metric_func(*inputs)
        if K.int_shape(metric)[-1] == 1:
            metric = K.squeeze(metric, axis=-1)
        if len(K.int_shape(metric)) == 0:
            metric = K.ones(K.shape(inputs[0])[0]) * metric

        # Apply the mask if needed
        if mask is not None:
            if not isinstance(mask, list):
                mask = [mask]
            mask = [K.cast(m, K.floatx()) for m in mask if m is not None]
            mask = reduce(lambda a, b: a*b, mask)
            metric *= mask
            metric /= K.mean(mask, axis=-1, keepdims=True)

        # Make sure that the tensor returned is (None, 1)
        dims = len(K.int_shape(metric))
        if dims > 1:
            metric = K.mean(metric, axis=list(range(1, dims)))

        return K.expand_dims(metric)

    @staticmethod
    def _generic_accuracy(y_true, y_pred):
        if K.int_shape(y_pred)[1] == 1:
            return binary_accuracy(y_true, y_pred)
        if K.int_shape(y_true)[-1] == 1:
            return sparse_categorical_accuracy(y_true, y_pred)

        return categorical_accuracy(y_true, y_pred)


class TripletLossLayer(Layer):
    """A bit of an unorthodox layer that implements the triplet loss with L2
    normalization.

    It receives 1 vector which is the concatenation of the three
    representations and performs the following operations.
    x = concat(x_a, x_p, x_n)
    N = x.shape[1]/3
    return ||x[:N] - x[2*N:]||_2^2 - ||x[:N] - x[N:2*N]||_2^2
    """
    def __init__(self, **kwargs):
        super(TripletLossLayer, self).__init__(**kwargs)

    def build(self, input_shape):
        assert not isinstance(input_shape, list)
        self.N = input_shape[1] // 3
        self.built = True

    def compute_output_shape(self, input_shape):
        assert not isinstance(input_shape, list)
        return (input_shape[0], 1)

    def call(self, x):
        N = self.N

        xa = x[:, :N]
        xp = x[:, N:2*N]
        xn = x[:, 2*N:]

        xa = xa / K.sqrt(K.sum(xa**2, axis=1, keepdims=True))
        xp = xp / K.sqrt(K.sum(xp**2, axis=1, keepdims=True))
        xn = xn / K.sqrt(K.sum(xn**2, axis=1, keepdims=True))

        dn = K.sum(K.square(xa - xn), axis=1, keepdims=True)
        dp = K.sum(K.square(xa - xp), axis=1, keepdims=True)

        return dn - dp
上述是metrics.py文件
上述为项目2
通过阅读项目1和项目2的代码可以很容易发现项目1并不是重要性采样，而项目2是重要性采样，同时项目1中引入了部分量子处理的内容，但是实际上他们的核心的训练过程都是类似的，无非是项目1采用了量子神经网络并且在采样的部分使用了量子优化器进行采样。项目1和项目2的采样部分都分为两层，第一层是随机采样，第二层是非均匀采样，项目1中采样的第二层采取的是使用量子优化器从N个样本中选取分数最大的K个数，项目2则是使样本根据分数来划分采样概率，分数越高的被采样的概率越高，从而实现了重要性采样并且是无偏的。而项目1由于无法使用概率来进行采样，并且没有设置P（x）和Q(x)，所以是有偏的。我现在需要你将项目1中的采样过程修改为无偏的，同时有如下要求：1.需要保证采样中的量子优化器部分不发生改变，这部分内容可以看作是封装的。2.在优化器封装的基础上添加内容使得其与项目2中类似可以产生概率，从而可以实现无偏采样。3.你需要仔细阅读项目2中的代码，并且明白其为什么是无偏的采样。我需要你按照类似的思想将项目1修改为无偏的采样。给出项目1修改后的对应代码