Keras自定义层与回调函数

import tensorflow as tf from tensorflow.keras import layers class CustomDense(layers.Layer): """自定义全连接层 演示 build() 创建权重、call() 定义前向传播 """ def __init__(self, units, activation=None, **kwargs): super().__init__(**kwargs) self.units = units self.activation = tf.keras.activations.get(activation) def build(self, input_shape): # input_shape: (batch_size, input_dim) input_dim = input_shape[-1] # 创建可训练权重: 核矩阵 W self.w = self.add_weight( shape=(input_dim, self.units), initializer='glorot_uniform', trainable=True, name='kernel' ) # 创建可训练偏置 b self.b = self.add_weight( shape=(self.units,), initializer='zeros', trainable=True, name='bias' ) super().build(input_shape) def call(self, inputs): # 前向传播: y = activation(x @ W + b) z = tf.matmul(inputs, self.w) + self.b if self.activation is not None: z = self.activation(z) return z def get_config(self): config = super().get_config() config.update({ 'units': self.units, 'activation': tf.keras.activations.serialize(self.activation) }) return config # 使用自定义层构建模型 model = tf.keras.Sequential([ CustomDense(64, activation='relu', input_shape=(128,)), CustomDense(32, activation='relu'), CustomDense(10, activation='softmax') ]) model.summary()

class CustomBatchNorm(layers.Layer): """自定义批归一化层 演示非可训练状态变量的管理 """ def __init__(self, momentum=0.99, epsilon=1e-3, **kwargs): super().__init__(**kwargs) self.momentum = momentum self.epsilon = epsilon def build(self, input_shape): channels = input_shape[-1] # 可训练的缩放因子 gamma 和偏移 beta self.gamma = self.add_weight( shape=(channels,), initializer='ones', trainable=True, name='gamma' ) self.beta = self.add_weight( shape=(channels,), initializer='zeros', trainable=True, name='beta' ) # 非可训练的状态变量：运行均值和方差 self.running_mean = self.add_weight( shape=(channels,), initializer='zeros', trainable=False, name='running_mean' ) self.running_variance = self.add_weight( shape=(channels,), initializer='ones', trainable=False, name='running_variance' ) super().build(input_shape) def call(self, inputs, training=None): if training: # 训练时：计算当前批次的均值和方差 batch_mean = tf.reduce_mean(inputs, axis=[0, 1, 2]) batch_var = tf.reduce_mean(tf.square(inputs - batch_mean), axis=[0, 1, 2]) # 更新运行均值和方差（指数移动平均） self.running_mean.assign( self.momentum * self.running_mean + (1 - self.momentum) * batch_mean ) self.running_variance.assign( self.momentum * self.running_variance + (1 - self.momentum) * batch_var ) mean, variance = batch_mean, batch_var else: # 推理时：使用运行均值和方差 mean, variance = self.running_mean, self.running_variance return tf.nn.batch_normalization( inputs, mean, variance, self.beta, self.gamma, self.epsilon )

class SelfAttention(layers.Layer): """简化版自注意力层（Scaled Dot-Product Attention） 输入形状: (batch_size, seq_len, embed_dim) 输出形状: (batch_size, seq_len, embed_dim) """ def __init__(self, embed_dim, num_heads=8, **kwargs): super().__init__(**kwargs) self.embed_dim = embed_dim self.num_heads = num_heads self.head_dim = embed_dim // num_heads assert self.head_dim * num_heads == embed_dim, \ "embed_dim must be divisible by num_heads" def build(self, input_shape): # 创建 Q, K, V 的投影矩阵 self.w_q = self.add_weight( shape=(self.embed_dim, self.embed_dim), initializer='glorot_uniform', trainable=True, name='w_q' ) self.w_k = self.add_weight( shape=(self.embed_dim, self.embed_dim), initializer='glorot_uniform', trainable=True, name='w_k' ) self.w_v = self.add_weight( shape=(self.embed_dim, self.embed_dim), initializer='glorot_uniform', trainable=True, name='w_v' ) self.w_out = self.add_weight( shape=(self.embed_dim, self.embed_dim), initializer='glorot_uniform', trainable=True, name='w_out' ) super().build(input_shape) def _split_heads(self, x, batch_size): # (batch, seq, embed) -> (batch, num_heads, seq, head_dim) x = tf.reshape(x, (batch_size, -1, self.num_heads, self.head_dim)) return tf.transpose(x, perm=[0, 2, 1, 3]) def call(self, inputs): batch_size = tf.shape(inputs)[0] # 线性投影 q = tf.matmul(tf.reshape(inputs, (-1, self.embed_dim)), self.w_q) k = tf.matmul(tf.reshape(inputs, (-1, self.embed_dim)), self.w_k) v = tf.matmul(tf.reshape(inputs, (-1, self.embed_dim)), self.w_v) # reshape 回 3D seq_len = tf.shape(inputs)[1] q = tf.reshape(q, (batch_size, seq_len, self.embed_dim)) k = tf.reshape(k, (batch_size, seq_len, self.embed_dim)) v = tf.reshape(v, (batch_size, seq_len, self.embed_dim)) # 分多头 q = self._split_heads(q, batch_size) # (b, h, s, d) k = self._split_heads(k, batch_size) v = self._split_heads(v, batch_size) # Scaled dot-product attention scale = tf.sqrt(tf.cast(self.head_dim, tf.float32)) scores = tf.matmul(q, k, transpose_b=True) / scale weights = tf.nn.softmax(scores, axis=-1) attn_output = tf.matmul(weights, v) # 合并多头 attn_output = tf.transpose(attn_output, perm=[0, 2, 1, 3]) attn_output = tf.reshape(attn_output, (batch_size, -1, self.embed_dim)) # 输出投影 attn_output = tf.matmul( tf.reshape(attn_output, (-1, self.embed_dim)), self.w_out ) attn_output = tf.reshape(attn_output, (batch_size, -1, self.embed_dim)) return attn_output def get_config(self): config = super().get_config() config.update({ 'embed_dim': self.embed_dim, 'num_heads': self.num_heads }) return config

class ReshapeFlatten(layers.Layer): """将任意形状展平为一维，并显式指定输出形状""" def compute_output_shape(self, input_shape): # input_shape: (batch, dim1, dim2, ...) batch = input_shape[0] total = 1 for d in input_shape[1:]: if d is None: total = None break total *= d return (batch, total) def call(self, inputs): batch_size = tf.shape(inputs)[0] return tf.reshape(inputs, (batch_size, -1))

class ResidualBlock(layers.Layer): """残差块：由两个卷积层和跳跃连接组成""" def __init__(self, filters, kernel_size=3, **kwargs): super().__init__(**kwargs) # 在 __init__ 中创建子层 self.conv1 = layers.Conv2D(filters, kernel_size, padding='same') self.bn1 = layers.BatchNormalization() self.conv2 = layers.Conv2D(filters, kernel_size, padding='same') self.bn2 = layers.BatchNormalization() self.relu = layers.ReLU() def call(self, inputs, training=None): shortcut = inputs x = self.conv1(inputs) x = self.bn1(x, training=training) x = self.relu(x) x = self.conv2(x) x = self.bn2(x, training=training) # 跳跃连接 x = layers.add([x, shortcut]) x = self.relu(x) return x

# ModelCheckpoint 完整配置 checkpoint = tf.keras.callbacks.ModelCheckpoint( filepath='model_epoch_{epoch:02d}_valacc_{val_accuracy:.4f}.h5', monitor='val_accuracy', # 监控验证集准确率 save_best_only=True, # 只保存最佳模型 save_weights_only=True, # 只保存权重（不保存完整模型） mode='max', # 监控指标越大越好 verbose=1 # 打印保存信息 ) # 也可以保存完整模型（包括优化器状态，支持恢复训练） checkpoint_full = tf.keras.callbacks.ModelCheckpoint( filepath='best_model.keras', monitor='val_loss', save_best_only=True, save_weights_only=False, # 保存完整模型 mode='min' ) model.fit( x_train, y_train, validation_data=(x_val, y_val), epochs=100, callbacks=[checkpoint, checkpoint_full] )

early_stop = tf.keras.callbacks.EarlyStopping( monitor='val_loss', # 监控验证损失 min_delta=1e-4, # 最小改善阈值 patience=10, # 容忍10个epoch无改善 verbose=1, mode='min', # val_loss 越小越好 restore_best_weights=True # 恢复最佳权重（关键！） ) # 与 ModelCheckpoint 配合使用 model.fit( x_train, y_train, validation_split=0.2, epochs=200, callbacks=[checkpoint, early_stop] ) # 当 EarlyStopping 触发时，自动恢复 val_loss 最低的权重

reduce_lr = tf.keras.callbacks.ReduceLROnPlateau( monitor='val_loss', factor=0.5, # 新学习率 = 旧学习率 * factor patience=5, # 5个epoch无改善则衰减 min_lr=1e-7, # 学习率下限 verbose=1, mode='min' ) # 查看学习率变化过程 class LRLogger(tf.keras.callbacks.Callback): def on_epoch_end(self, epoch, logs=None): lr = self.model.optimizer.lr.numpy() print(f"Epoch {epoch+1}: learning rate = {lr:.2e}") model.fit( x_train, y_train, validation_split=0.2, epochs=100, callbacks=[reduce_lr, LRLogger()] )

# 余弦退火学习率调度 def cosine_annealing(epoch, lr): """余弦退火: 从初始学习率按余弦曲线下降到0""" initial_lr = 1e-3 total_epochs = 100 return initial_lr * 0.5 * (1 + tf.cos(np.pi * epoch / total_epochs)) # 分段常数衰减 def step_decay(epoch, lr): """每30个epoch学习率减半""" if epoch < 30: return 1e-3 elif epoch < 60: return 5e-4 elif epoch < 90: return 1e-4 else: return 5e-5 lr_scheduler = tf.keras.callbacks.LearningRateScheduler( schedule=cosine_annealing, verbose=1 )

csv_logger = tf.keras.callbacks.CSVLogger( filename='training_log.csv', separator=',', append=False # False: 覆盖写入; True: 追加写入 ) # 读取CSV日志进行分析 import pandas as pd df = pd.read_csv('training_log.csv') # 绘制训练曲线 plt.plot(df['epoch'], df['loss'], label='train_loss') plt.plot(df['epoch'], df['val_loss'], label='val_loss') plt.legend() plt.show()

class TrainingMonitor(tf.keras.callbacks.Callback): """训练监控回调：实时追踪梯度范数、权重更新量等内部状态""" def __init__(self, log_freq=10): super().__init__() self.log_freq = log_freq def on_train_begin(self, logs=None): self.grad_norms = [] self.weight_updates = [] print("=" * 50) print("开始训练监控...") print(f"模型层数: {len(self.model.layers)}") for i, layer in enumerate(self.model.layers): if hasattr(layer, 'kernel'): print(f" 层 {i}: {layer.name}, 核形状: {layer.kernel.shape}") def on_batch_end(self, batch, logs=None): if batch % self.log_freq == 0: loss = logs.get('loss') acc = logs.get('accuracy') print(f" Batch {batch:5d}: loss={loss:.4f}, acc={acc:.4f}") def on_epoch_end(self, epoch, logs=None): val_loss = logs.get('val_loss') val_acc = logs.get('val_accuracy') lr = self.model.optimizer.lr.numpy() print(f">>> Epoch {epoch+1}: val_loss={val_loss:.4f}, " f"val_acc={val_acc:.4f}, lr={lr:.2e}")

class CoolDownEarlyStopping(tf.keras.callbacks.Callback): """带冷却周期的早停回调 检测到指标停滞 -> 降低学习率（冷却）-> 若继续停滞 -> 停止训练 """ def __init__(self, monitor='val_loss', patience=10, cooldown_patience=5, cooldown_factor=0.5, min_delta=1e-4, min_lr=1e-7): super().__init__() self.monitor = monitor self.patience = patience self.cooldown_patience = cooldown_patience self.cooldown_factor = cooldown_factor self.min_delta = min_delta self.min_lr = min_lr self.reset() def reset(self): self.wait = 0 self.cooldown_counter = 0 self.best_weights = None self.best = float('inf') self.stopped_epoch = 0 def on_epoch_end(self, epoch, logs=None): current = logs.get(self.monitor) if current is None: return if current < self.best - self.min_delta: # 指标改善：重置等待计数、保存最佳权重 self.best = current self.wait = 0 self.cooldown_counter = 0 self.best_weights = self.model.get_weights() else: self.wait += 1 if self.wait >= self.patience: # 进入冷却期，降低学习率 if self.cooldown_counter < self.cooldown_patience: old_lr = self.model.optimizer.lr.numpy() new_lr = old_lr * self.cooldown_factor if new_lr > self.min_lr: self.model.optimizer.lr.assign(new_lr) print(f">>> CoolDown: lr {old_lr:.2e} -> {new_lr:.2e}") self.cooldown_counter += 1 self.wait = 0 # 重置等待计数 else: # 冷却期已过，停止训练 self.stopped_epoch = epoch self.model.stop_training = True print(f">>> Early stopping at epoch {epoch+1}") # 每5个epoch打印一次状态 if epoch % 5 == 0: print(f" [{self.monitor}] best={self.best:.4f}, " f"wait={self.wait}/{self.patience}, " f"cooldown={self.cooldown_counter}/{self.cooldown_patience}") def on_train_end(self, logs=None): # 恢复最佳权重 if self.best_weights is not None: self.model.set_weights(self.best_weights) print("恢复最佳模型权重")

def train_with_advanced_callbacks(model, x_train, y_train, x_val, y_val, epochs=200): """训练函数：组合多种回调实现智能训练""" # 1. 模型保存 checkpoint = tf.keras.callbacks.ModelCheckpoint( filepath='best_model.h5', monitor='val_accuracy', save_best_only=True, mode='max' ) # 2. 早停 early_stop = tf.keras.callbacks.EarlyStopping( monitor='val_loss', patience=20, restore_best_weights=True ) # 3. 自适应学习率 reduce_lr = tf.keras.callbacks.ReduceLROnPlateau( monitor='val_loss', factor=0.5, patience=5, min_lr=1e-6 ) # 4. 自定义进度显示 class ProgressBar(tf.keras.callbacks.Callback): def __init__(self, total_epochs): super().__init__() self.total_epochs = total_epochs self.start_time = None def on_train_begin(self, logs=None): import time self.start_time = time.time() print(f"开始训练，共 {self.total_epochs} 个epochs") def on_epoch_end(self, epoch, logs=None): import time elapsed = time.time() - self.start_time pct = (epoch + 1) / self.total_epochs * 100 bar = '#' * int(pct // 5) + '-' * (20 - int(pct // 5)) print(f"\r[{bar}] {pct:.0f}% | " f"loss={logs['loss']:.4f} | " f"val_acc={logs['val_accuracy']:.4f} | " f"time={elapsed:.0f}s", end='') # 5. 日志记录 csv_logger = tf.keras.callbacks.CSVLogger('training_log.csv') # 组合所有回调 callbacks = [ checkpoint, early_stop, reduce_lr, ProgressBar(epochs), csv_logger ] # 开始训练 history = model.fit( x_train, y_train, validation_data=(x_val, y_val), epochs=epochs, callbacks=callbacks, verbose=0 # 关闭Keras默认进度条 ) return history

# 基础配置 tensorboard = tf.keras.callbacks.TensorBoard( log_dir='./logs/fit/', # 日志目录 histogram_freq=1, # 每N个epoch记录权重直方图 write_graph=True, # 记录计算图 write_images=True, # 记录权重可视化图像 update_freq='epoch', # 记录频率：'epoch' | 'batch' | 整数(全局步数) profile_batch=2 # 分析第2个batch的性能 ) model.fit( x_train, y_train, epochs=50, validation_data=(x_val, y_val), callbacks=[tensorboard] ) # 启动 TensorBoard # 命令行: tensorboard --logdir ./logs/fit # 浏览器打开: http://localhost:6006

import io import matplotlib.pyplot as plt class CustomMetricsLogger(tf.keras.callbacks.Callback): """通过 tf.summary 记录自定义指标到 TensorBoard""" def __init__(self, log_dir='./logs/custom'): super().__init__() self.file_writer = tf.summary.create_file_writer(log_dir) def on_epoch_end(self, epoch, logs=None): with self.file_writer.as_default(): # 记录所有标准指标 for name, value in logs.items(): tf.summary.scalar(name, value, step=epoch) # 记录自定义指标：学习率 lr = self.model.optimizer.lr.numpy() tf.summary.scalar('learning_rate', lr, step=epoch) # 记录权重统计信息 for i, layer in enumerate(self.model.layers): weights = layer.get_weights() for j, w in enumerate(weights): tf.summary.histogram( f'layer_{i}/{layer.name}_w{j}', w, step=epoch ) tf.summary.scalar( f'layer_{i}/{layer.name}_w{j}_norm', tf.norm(w).numpy(), step=epoch ) # 使用 custom_logger = CustomMetricsLogger('./logs/custom') model.fit(x_train, y_train, epochs=50, callbacks=[custom_logger])

class ImageLogger(tf.keras.callbacks.Callback): """在TensorBoard中记录预测结果图像""" def __init__(self, x_test, y_test, class_names, log_dir='./logs/images'): super().__init__() self.x_test = x_test self.y_test = y_test self.class_names = class_names self.file_writer = tf.summary.create_file_writer(log_dir) def on_epoch_end(self, epoch, logs=None): # 每10个epoch记录一次 if epoch % 10 != 0: return # 随机选取8个测试样本 indices = np.random.choice(len(self.x_test), 8) x_batch = self.x_test[indices] y_batch = self.y_test[indices] # 预测 preds = self.model.predict(x_batch, verbose=0) pred_classes = np.argmax(preds, axis=1) # 创建图像 fig, axes = plt.subplots(2, 4, figsize=(12, 6)) for i, ax in enumerate(axes.flat): ax.imshow(x_batch[i]) true_label = self.class_names[y_batch[i]] pred_label = self.class_names[pred_classes[i]] color = 'green' if y_batch[i] == pred_classes[i] else 'red' ax.set_title(f"True: {true_label}\nPred: {pred_label}", color=color, fontsize=10) ax.axis('off') plt.tight_layout() # 将 matplotlib 图像转换为 TensorBoard 可读格式 buf = io.BytesIO() plt.savefig(buf, format='png', dpi=100) plt.close() buf.seek(0) image = tf.image.decode_png(buf.getvalue(), channels=4) image = tf.expand_dims(image, 0) with self.file_writer.as_default(): tf.summary.image('predictions', image, step=epoch)

import tensorboard.plugins.projector as projector def log_embeddings(model, data, labels, label_names, log_dir): """记录嵌入向量到TensorBoard Embedding Projector""" # 获取嵌入层的输出 embed_model = tf.keras.Model( inputs=model.input, outputs=model.get_layer('embedding').output ) embeddings = embed_model.predict(data) # 保存权重文件 weights_path = os.path.join(log_dir, 'embedding.ckpt') checkpoint = tf.train.Checkpoint(embedding=tf.Variable(embeddings)) checkpoint.save(weights_path) # 保存元数据（标签） metadata_path = os.path.join(log_dir, 'metadata.tsv') with open(metadata_path, 'w') as f: f.write('label\tclass\n') for label in labels: f.write(f"{label_names[label]}\t{label}\n") # 配置 projector config = projector.ProjectorConfig() embedding = config.embeddings.add() embedding.tensor_name = 'embedding/.ATTRIBUTES/VARIABLE_VALUE' embedding.metadata_path = 'metadata.tsv' projector.visualize_embeddings(log_dir, config)

from tensorboard.plugins.hparams import api as hp # 定义超参数空间 HP_NUM_UNITS = hp.HParam('num_units', hp.Discrete([32, 64, 128])) HP_DROPOUT = hp.HParam('dropout', hp.RealInterval(0.1, 0.5)) HP_OPTIMIZER = hp.HParam('optimizer', hp.Discrete(['adam', 'sgd'])) HP_LR = hp.HParam('learning_rate', hp.Discrete([1e-3, 1e-4])) METRIC_ACCURACY = hp.Metric( 'accuracy', display_name='Validation Accuracy' ) def build_and_train(hparams, x_train, y_train, x_val, y_val, log_dir): """使用指定超参数构建并训练模型""" model = tf.keras.Sequential([ layers.Flatten(input_shape=(28, 28)), layers.Dense(hparams[HP_NUM_UNITS], activation='relu'), layers.Dropout(hparams[HP_DROPOUT]), layers.Dense(10, activation='softmax') ]) if hparams[HP_OPTIMIZER] == 'adam': optimizer = tf.keras.optimizers.Adam(hparams[HP_LR]) else: optimizer = tf.keras.optimizers.SGD(hparams[HP_LR]) model.compile( optimizer=optimizer, loss='sparse_categorical_crossentropy', metrics=['accuracy'] ) # HParams回调：自动记录超参数和结果 hparams_callback = hp.KerasCallback(log_dir, hparams) model.fit( x_train, y_train, validation_data=(x_val, y_val), epochs=10, callbacks=[hparams_callback], verbose=0 ) return model # 超参数搜索循环 session_num = 0 for num_units in HP_NUM_UNITS.domain.values: for dropout_rate in np.arange(0.1, 0.51, 0.2): for optimizer in HP_OPTIMIZER.domain.values: for lr in HP_LR.domain.values: hparams = { HP_NUM_UNITS: num_units, HP_DROPOUT: dropout_rate, HP_OPTIMIZER: optimizer, HP_LR: lr } run_name = f"run_{session_num}" print(f"--- Running {run_name}: {hparams}") log_dir = f"./logs/hparam_tuning/{run_name}" build_and_train(hparams, x_train, y_train, x_val, y_val, log_dir) session_num += 1 # 启动 TensorBoard 查看 HParams 仪表盘 # tensorboard --logdir ./logs/hparam_tuning

""" 综合实战：使用自定义层 + 高级回调构建图像分类器 包含：自定义卷积注意力层 + TensorBoard + 早停 + 模型检查点 """ import tensorflow as tf from tensorflow.keras import layers, models import numpy as np # ====== 1. 自定义卷积注意力层 ====== class ConvAttention(layers.Layer): """卷积注意力模块：通道注意力 + 空间注意力""" def __init__(self, reduction_ratio=16, **kwargs): super().__init__(**kwargs) self.reduction_ratio = reduction_ratio def build(self, input_shape): channels = input_shape[-1] # 通道注意力网络 self.avg_pool = layers.GlobalAveragePooling2D() self.max_pool = layers.GlobalMaxPooling2D() self.dense1 = layers.Dense(channels // self.reduction_ratio, activation='relu') self.dense2 = layers.Dense(channels, activation='sigmoid') # 空间注意力卷积 self.spatial_conv = layers.Conv2D( 1, kernel_size=7, padding='same', activation='sigmoid' ) super().build(input_shape) def call(self, inputs): # 通道注意力分支 avg_out = self.dense2(self.dense1(self.avg_pool(inputs))) max_out = self.dense2(self.dense1(self.max_pool(inputs))) channel_att = tf.expand_dims(tf.expand_dims( avg_out + max_out, 1), 1) # 空间注意力分支 avg_spatial = tf.reduce_mean(inputs, axis=-1, keepdims=True) max_spatial = tf.reduce_max(inputs, axis=-1, keepdims=True) spatial_concat = tf.concat([avg_spatial, max_spatial], axis=-1) spatial_att = self.spatial_conv(spatial_concat) # 应用注意力 return inputs * channel_att * spatial_att # ====== 2. 自定义学习率调度回调 ====== class WarmUpCosineDecay(tf.keras.callbacks.Callback): """学习率先预热后余弦衰减""" def __init__(self, warmup_epochs=5, initial_lr=1e-4, max_lr=1e-3, total_epochs=100): super().__init__() self.warmup_epochs = warmup_epochs self.initial_lr = initial_lr self.max_lr = max_lr self.total_epochs = total_epochs def on_epoch_begin(self, epoch, logs=None): if epoch < self.warmup_epochs: # 线性预热 lr = self.initial_lr + (self.max_lr - self.initial_lr) * \ (epoch / self.warmup_epochs) else: # 余弦衰减 progress = (epoch - self.warmup_epochs) / \ (self.total_epochs - self.warmup_epochs) lr = self.max_lr * 0.5 * (1 + np.cos(np.pi * progress)) self.model.optimizer.lr.assign(lr) # ====== 3. 构建模型 ====== def build_attention_model(input_shape=(28, 28, 1), num_classes=10): inputs = layers.Input(shape=input_shape) x = layers.Conv2D(32, 3, padding='same', activation='relu')(inputs) x = ConvAttention()(x) x = layers.MaxPooling2D(2)(x) x = layers.Conv2D(64, 3, padding='same', activation='relu')(x) x = ConvAttention()(x) x = layers.MaxPooling2D(2)(x) x = layers.GlobalAveragePooling2D()(x) x = layers.Dense(128, activation='relu')(x) x = layers.Dropout(0.5)(x) outputs = layers.Dense(num_classes, activation='softmax')(x) return models.Model(inputs, outputs) # ====== 4. 数据准备与训练 ====== (x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data() x_train = x_train.astype('float32')[..., np.newaxis] / 255.0 x_test = x_test.astype('float32')[..., np.newaxis] / 255.0 model = build_attention_model() model.compile( optimizer=tf.keras.optimizers.Adam(1e-3), loss='sparse_categorical_crossentropy', metrics=['accuracy'] ) # ====== 5. 配置回调 ====== callbacks = [ tf.keras.callbacks.ModelCheckpoint( 'best_attention_model.h5', monitor='val_accuracy', save_best_only=True ), tf.keras.callbacks.EarlyStopping( monitor='val_loss', patience=15, restore_best_weights=True ), tf.keras.callbacks.ReduceLROnPlateau( monitor='val_loss', factor=0.5, patience=5 ), tf.keras.callbacks.TensorBoard( log_dir='./logs/attention_mnist', histogram_freq=1 ), WarmUpCosineDecay( warmup_epochs=3, initial_lr=1e-5, max_lr=1e-3, total_epochs=50 ) ] # 训练 history = model.fit( x_train, y_train, validation_split=0.2, epochs=50, batch_size=128, callbacks=callbacks, verbose=1 ) # 评估 test_loss, test_acc = model.evaluate(x_test, y_test) print(f"测试准确率: {test_acc:.4f}") print(f"测试损失: {test_loss:.4f}")

# 保存模型 model.save('my_model.keras') # 加载模型时需要指定自定义层 model = tf.keras.models.load_model( 'my_model.keras', custom_objects={'CustomDense': CustomDense, 'SelfAttention': SelfAttention} )