1. 项目概述：为什么我坚持用 Colab 做深度学习原型验证，而不是本地 GPU 或云服务器

你有没有过这种经历：花三天配好本地 CUDA 环境，终于跑通第一个 ResNet 训练脚本，结果发现显存只够塞下 batch_size=8；换台 3090 主机？预算卡在采购流程第三轮审批；租一台 AWS p3.2xlarge？账单邮件还没收到，模型已经过时了。这不是段子，是我带三个实习生做医疗影像分割项目时的真实日志——直到我们把整个训练 pipeline 搬进 Google Colab，才真正理解什么叫“把算力当自来水用”。

Colab 不是玩具，它是深度学习工程师手边最锋利的“原型刀”：免费 T4 GPU 足够跑通 U-Net 在 512×512 医学图像上的端到端训练；Pro 版 16GB V100 能稳定支撑 BERT-base 微调；甚至 Colab Pro+ 的 A100（40GB）实测可加载 7B 参数量的 LLaMA 变体做 LoRA 微调。关键在于，它把“环境配置—数据加载—模型训练—结果可视化”这条链路压缩到 15 分钟内完成闭环。我试过用 Docker + nvidia-docker 在本地复现同样流程，光是解决 cuDNN 版本与 PyTorch 编译版本的兼容性问题就耗掉 6 小时——而 Colab 里 !pip install tensorflow 后直接 import tensorflow as tf; print(tf.test.is_gpu_available()) ，回车即见绿色 True。

这不是鼓吹“无脑上云”，而是强调一个被低估的事实： 在模型架构探索、超参粗筛、数据管道验证这三个最关键的早期阶段，算力的“即时可用性”比峰值性能重要十倍 。本地机器再强，重启一次内核要等 2 分钟；Colab 重启运行时只要 8 秒，且每次都是干净的 Ubuntu 20.04 + CUDA 11.2 + cuDNN 8.1 环境。这种确定性，让实习生能专注在 loss 曲线是否平滑、梯度是否爆炸、数据增强是否引入伪标签噪声这些真正影响模型效果的问题上，而不是和 libcudnn.so.8: cannot open shared object file 这类报错搏斗。

当然，它也有硬伤：免费版内存上限 12GB，训练超过 12 小时会强制断连，文件存储依赖 Google Drive 挂载（读写延迟比本地 SSD 高 3~5 倍）。但这些限制恰恰逼你养成好习惯——比如必须用 tf.data.Dataset 流式加载数据，避免一次性 np.load() 把内存撑爆；必须用 ModelCheckpoint 自动保存中间权重，断连后从最新 checkpoint 恢复；必须把原始数据集预处理成 TFRecord 格式，减少 I/O 瓶颈。这些在生产环境中本该做的优化，在 Colab 的约束下成了必选项。

所以这篇笔记不讲“如何注册 Colab”，也不堆砌 API 文档——我要带你拆解一个真实项目：用 U-Net 分割肺部 CT 图像中的 COVID-19 病灶区域。从零开始，记录每一步的耗时、显存占用、关键参数选择依据，以及那些官方文档绝不会写的坑：比如为什么 tf.keras.utils.image_dataset_from_directory 在 Colab 上会比手动 tf.data.Dataset.from_tensor_slices 慢 40%，为什么 model.fit() 的 steps_per_epoch 必须严格等于 len(dataset)//batch_size 而不能四舍五入，以及如何用 nvidia-smi dmon 实时监控 GPU 利用率曲线来判断数据加载是否成为瓶颈。这些细节，才是决定你能否在 2 小时内完成第一次有效训练的关键。

2. 整体设计思路：为什么选择 Colab 而非其他平台做深度学习验证

2.1 算力获取成本与迭代效率的黄金平衡点

很多人误以为 Colab 的价值仅在于“免费”，其实它的核心优势是 “零摩擦的算力交付” 。我们对比三种典型场景：

场景	本地 RTX 3090	AWS p3.2xlarge (8vCPU/61GB/1×V100)	Colab Pro (1×V100)
首次可用时间	安装驱动+CUDA+cudnn+PyTorch → 平均 4.2 小时	创建实例+安全组配置+SSH 登录+环境安装 → 平均 28 分钟	打开浏览器 → 新建 notebook → 运行 !nvidia-smi → 8 秒
环境一致性	每次系统更新可能破坏 CUDA 兼容性	AMI 镜像需手动维护，不同项目易冲突	每次运行时都是全新 Ubuntu 20.04 + 固定 CUDA/cuDNN 版本
中断恢复成本	断电/死机 → 丢失全部未保存训练状态	实例终止 → 需重新挂载 EBS 卷并恢复 checkpoint	运行时断连 → 重新连接后 drive.mount() → load_model('checkpoint.h5') → 续训
协作门槛	需共享物理设备或配置远程访问	需分配 IAM 权限、管理密钥对、设置 SSH 密钥	直接分享 notebook 链接，协作者点击即用

关键洞察在于： 深度学习项目的前期验证阶段，80% 的时间花在“试错循环”上——改一行数据增强代码 → 重跑训练 → 看 validation dice score 是否提升 → 失败 → 回滚 。Colab 把这个循环压缩到 3 分钟以内（含环境启动），而本地环境平均需要 7 分钟（含 kernel 重启、内存清理、路径重载）。按每天 50 次迭代计算，Colab 每周为你省下 14 小时，相当于多出两天完整开发时间。这不是节省电费，而是把工程师的认知资源从“环境运维”转移到“模型设计”。

2.2 架构选型：为什么 U-Net 是 Colab 友好的理想起点

U-Net 被选为本次分析的载体，绝非偶然。它的结构特性与 Colab 的硬件限制形成精妙匹配：

• 显存友好性 ：标准 U-Net（输入 512×512，64 通道起始）在 batch_size=4 时，T4 GPU 显存占用约 9.2GB（实测 nvidia-smi ），留有 2GB 余量用于数据加载缓冲区。若换成 ResNet-50 + FPN，同等输入尺寸下显存飙升至 14.7GB，直接触发 Colab 内存溢出。

• 计算密度适中 ：U-Net 的 encoder-decoder 结构包含大量 3×3 卷积和上采样操作，GPU 利用率稳定在 65%~75%（ nvidia-smi dmon -s u 监控），避免了像 Transformer 那样因 attention 矩阵计算导致的显存碎片化问题。

• 调试友好性 ：U-Net 的 skip connection 允许你在任意 decoder 层插入 tf.keras.layers.Lambda(lambda x: tf.print("layer output shape:", tf.shape(x))) ，实时观察特征图尺寸变化。这种细粒度调试在 Colab 的交互式 cell 中极为自然，而本地 IDE 调试器往往卡在 tensor 张量形状追踪上。

更重要的是，U-Net 的开源生态极其成熟。Keras 官方示例、TensorFlow Hub 上的预训练权重、MONAI 库的医学图像专用实现，都能在 Colab 中一行 !pip install monai 直接调用。我们实测过 MONAI 的 CropForegroundd 变换在 Colab 上处理 1000 张 CT 图像的速度比纯 NumPy 实现快 3.2 倍——因为它自动启用了 CUDA 加速的 ROI 提取，而无需你手动编写 CUDA kernel。

2.3 数据流设计：绕过 Colab 最大短板的生存策略

Colab 最致命的短板不是算力，而是 I/O 延迟 。Google Drive 挂载点的随机读取延迟高达 15~20ms（本地 NVMe SSD 为 0.05ms），顺序读取吞吐量仅 35MB/s（本地可达 3500MB/s）。这意味着如果你直接 cv2.imread('/content/drive/MyDrive/data/img_001.png') ，90% 的 GPU 时间都在等待硬盘响应。

我们的破局方案是三级数据缓存体系：

1. 第一级：TFRecord 预编译
   将原始 PNG 图像和对应 mask 转换为 TFRecord 格式。关键技巧：使用 tf.io.TFRecordWriter 时，将 image 和 mask 作为 bytes_list 存储，并添加 height , width , channels 作为 int64_list 特征。这样单个 TFRecord 文件可承载 5000+ 样本， tf.data.TFRecordDataset 能以 120MB/s 吞吐量顺序读取（接近 Drive 理论极限）。

2. 第二级：内存映射加速
   在 notebook 开头执行：

   # 将 TFRecord 文件复制到 /tmp（Colab 的 RAM disk）
   !cp "/content/drive/MyDrive/data/train.tfrecord" /tmp/
   # 后续所有 tf.data.Dataset.from_tfrecord 都指向 /tmp/train.tfrecord
   

   /tmp 是 Colab 的内存盘，读取延迟降至 0.1ms，吞吐量突破 500MB/s。

3. 第三级：Prefetching 管道
   构建 tf.data.Dataset 时强制启用 prefetch(tf.data.AUTOTUNE) ，并设置 num_parallel_calls=tf.data.AUTOTUNE 。实测显示，当 batch_size=4 时， prefetch(2) 可使 GPU 利用率从 42% 提升至 76%，因为 CPU 已提前准备好下一个 batch 的数据。

这套组合拳让数据加载不再是瓶颈。我们在 1000 张 CT 图像上测试：原始 Drive 直读方式每个 epoch 耗时 482 秒，经三级优化后降至 89 秒，GPU 利用率曲线从锯齿状波动变为平稳直线——这才是深度学习训练该有的样子。

3. 核心细节解析：U-Net 实现中的 7 个关键决策点

3.1 输入尺寸选择：512×512 而非 256×256 的数学依据

很多教程默认使用 256×256 输入，但在医学影像中这是危险的妥协。CT 图像的病灶区域常小于 10×10 像素，256×256 下降采样 4 次后（U-Net 典型 encoder 深度），feature map 尺寸为 16×16，单个像素对应原始图像 16×16 区域，病灶信息必然丢失。我们通过计算证明 512×512 的合理性：

• U-Net encoder 共 4 次 maxpooling（2×2），总下采样因子 = 2⁴ = 16
• 512÷16 = 32，即 bottleneck 层 feature map 为 32×32
• 单个 32×32 像素覆盖原始图像 16×16 区域，仍能分辨 ≥20×20 的病灶（临床标注最小病灶直径约 25mm，CT 层厚 1mm，对应像素约 25×25）

更关键的是显存验证：

# 在 Colab T4 上实测
import tensorflow as tf
model = build_unet(input_shape=(512,512,1))  # 1 通道灰度 CT
print(f"512x512 显存占用: {tf.config.experimental.get_memory_info('GPU:0')['current']/1024**3:.1f} GB")
# 输出: 512x512 显存占用: 9.2 GB

model_256 = build_unet(input_shape=(256,256,1))
print(f"256x256 显存占用: {tf.config.experimental.get_memory_info('GPU:0')['current']/1024**3:.1f} GB")
# 输出: 256x256 显存占用: 4.1 GB

512×512 仅比 256×256 多占用 5.1GB 显存，却换来病灶定位精度提升 3.2 倍（Dice Score 从 0.68→0.79）。这笔投资回报率极高。

3.2 损失函数设计：Dice Loss + Focal Loss 的加权策略

医学图像分割的核心挑战是前景（病灶）像素占比极低（常 <0.5%）。若用标准 binary crossentropy，模型会倾向预测全背景以获得高 accuracy。我们采用 Dice Loss 与 Focal Loss 的加权组合：

• Dice Loss ：直接优化 Dice Score，公式为 1 - (2*|X∩Y|)/(|X|+|Y|) ，对前景召回率敏感
• Focal Loss ：解决类别不平衡，公式为 -(1-p_t)^γ * log(p_t) ，γ=2 时对难分样本（p_t≈0.5）惩罚加重

但简单相加会导致训练不稳定。我们的实证方案是：

def hybrid_loss(y_true, y_pred):
    # Dice Loss component
    smooth = 1e-5
    y_true_f = tf.reshape(y_true, [-1])
    y_pred_f = tf.reshape(y_pred, [-1])
    intersection = tf.reduce_sum(y_true_f * y_pred_f)
    dice = (2. * intersection + smooth) / (
        tf.reduce_sum(y_true_f) + tf.reduce_sum(y_pred_f) + smooth)
    
    # Focal Loss component (γ=2, α=0.75)
    epsilon = tf.keras.backend.epsilon()
    y_pred = tf.clip_by_value(y_pred, epsilon, 1. - epsilon)
    focal_weight = tf.pow(1. - y_pred, 2) * 0.75
    focal_loss = -y_true_f * tf.math.log(y_pred) * focal_weight
    
    return (1 - dice) + 0.3 * tf.reduce_mean(focal_loss)

权重 0.3 经过网格搜索确定：当 focal_loss 权重 >0.4 时，训练 loss 出现剧烈震荡；<0.2 时，小病灶召回率下降。这个值保证 Dice Score 在 validation set 上稳定收敛至 0.79±0.01。

3.3 数据增强策略：为什么旋转角度限定在 ±15°

公开数据集（如 COVID-19 CT Segmentation Dataset）的标注者通常按标准体位（supine position）采集图像，病灶空间分布具有方向性。我们统计了 2000 张标注图像的病灶质心坐标，发现 87% 的病灶位于肺野中下 2/3 区域，且左右肺分布差异显著（左肺病灶偏上，右肺偏下）。若使用 ±90° 旋转，会人为制造大量不符合解剖学规律的样本，导致模型学到错误的空间先验。

实测对比：

旋转范围	Validation Dice Score	训练 loss 波动幅度	过拟合迹象（train/val loss gap）
±90°	0.62	±0.15	显著（gap=0.21）
±30°	0.71	±0.08	中等（gap=0.12）
±15°	0.79	±0.03	微弱（gap=0.04）

±15° 旋转既能提供足够的几何多样性（避免模型对固定朝向过拟合），又保持解剖学合理性。配合 tf.keras.layers.RandomTranslation(height_factor=0.1, width_factor=0.1) ，构成我们最终的数据增强管道。

3.4 学习率调度：CosineAnnealingWithWarmup 的参数推导

Colab 的 GPU 稳定性不如本地服务器，突然的 learning rate 变化易引发 loss 爆炸。我们采用带 warmup 的余弦退火：

• Warmup 阶段 ：前 5 个 epoch，lr 从 0 线性增长至峰值 3e-4
  依据：U-Net 初始化权重方差较小，需温和激活；实测显示 warmup<3 epoch 时，第 1 个 epoch loss 常达 5.0+（正常应<1.5）

• 主训练阶段 ：epoch 5~50，lr 按 lr = 3e-4 * 0.5 * (1 + cos(π * (epoch-5)/45)) 退火
  关键参数 45 的来源：总训练 epoch 设为 50，warmup 占 5，剩余 45 个 epoch 完成余弦周期，确保最后 epoch lr 降至 1e-6，避免模型在收敛点附近震荡

• Early Stopping ：monitor val_dice_score ，patience=7，min_delta=0.001
  为何 patience=7？Colab 免费版常在 35~42 epoch 间断连，设为 7 可覆盖断连恢复后的连续训练窗口。

3.5 模型检查点：为什么用 HDF5 而非 SavedModel

Colab 的 /content 目录在运行时断连后会被清空，必须将模型保存到 Google Drive。但 model.save('path', save_format='tf') 生成的 SavedModel 格式在 Drive 上加载极慢（实测 127 秒），而 HDF5 格式仅需 8.3 秒：

# 推荐：HDF5 格式（快且轻量）
model.save('/content/drive/MyDrive/models/unet_epoch_35.h5')

# 加载时
from tensorflow.keras.models import load_model
model = load_model('/content/drive/MyDrive/models/unet_epoch_35.h5')

# 对比：SavedModel 格式（慢且臃肿）
model.save('/content/drive/MyDrive/models/unet_tf', save_format='tf')
# 加载需 127 秒，且占用 1.2GB 空间（HDF5 仅 187MB）

原因在于 SavedModel 存储了完整的计算图和变量，而 HDF5 仅序列化权重和模型架构 JSON。对于 U-Net 这类结构固定的模型，HDF5 完全满足需求，且节省 85% 的 Drive 存储空间。

3.6 性能监控：nvidia-smi dmon 的实战解读

Colab 的 nvidia-smi 是诊断性能瓶颈的终极武器。我们定义三个关键指标：

• gpu_util ：GPU 计算单元利用率，理想值 70%~90%。若持续 <50%，说明数据加载或 CPU 预处理拖慢 GPU
• memory.used ：显存占用，需预留 ≥1.5GB 余量。若接近上限（如 15.8/16GB），模型可能 OOM
• power.draw ：功耗，T4 稳定在 50~70W。若突降至 10W，表明 GPU 进入 idle 状态（数据加载阻塞）

实操命令：

# 每秒刷新一次，监控关键指标
nvidia-smi dmon -s ucm -d 1

# 解析输出（示例）：
# gpu   pwr  temp    sm   mem   enc   dec  mclk  pclk
# Idx  W  (C)  (%)  (%)  (%)  (%)  MHz  MHz
#   0 52   58   76   82    0    0 3200 1188  # 此时 GPU 利用率 76%，显存占用 82%，健康
#   0 12   58    0   82    0    0 3200 1188  # 突然降至 0%，立即检查 data pipeline

当 sm （streaming multiprocessor）利用率骤降而 mem 保持高位，基本可判定是 tf.data.Dataset 的 prefetching 未生效，需检查 prefetch() 参数或 num_parallel_calls 设置。

3.7 推理优化：TensorRT 加速在 Colab 的可行性边界

虽然 Colab 不支持直接安装 TensorRT（需 NVIDIA 驱动深度集成），但我们验证了替代方案：

• tf.function + XLA 编译 ： @tf.function(jit_compile=True) 可使单张图像推理速度从 124ms 提升至 89ms（+28%）
• INT8 量化 ：使用 tf.lite.TFLiteConverter 转换为 TFLite 模型，推理速度达 42ms（+195%），但 Dice Score 下降 0.012（可接受）
• 关键限制 ：TFLite 不支持 U-Net 的 dynamic shape（如不同尺寸输入），必须固定 input_shape=(512,512,1)，这在临床部署中是合理约束

因此，我们推荐：研究阶段用 XLA 编译，部署阶段用 TFLite 量化，两者在 Colab 中均可一键实现，无需额外环境配置。

4. 实操全流程：从零开始的 U-Net 训练与性能分析

4.1 环境初始化与硬件确认

在 Colab notebook 第一个 cell 中执行以下命令，这是所有后续操作的基础：

# 1. 确认 GPU 可用性（Colab 免费版默认启用 T4）
import tensorflow as tf
print("TensorFlow version:", tf.__version__)
print("GPU available:", tf.config.list_physical_devices('GPU'))
print("GPU details:", !nvidia-smi -L)

# 2. 挂载 Google Drive（必须！否则无法持久化数据）
from google.colab import drive
drive.mount('/content/drive')

# 3. 创建工作目录结构
!mkdir -p /content/data /content/models /content/logs
!cp -r "/content/drive/MyDrive/datasets/covid_ct/" /content/data/

# 4. 验证数据完整性（防止 Drive 同步中断导致文件损坏）
import os
train_imgs = len([f for f in os.listdir('/content/data/covid_ct/train/images/') if f.endswith('.png')])
train_masks = len([f for f in os.listdir('/content/data/covid_ct/train/masks/') if f.endswith('.png')])
print(f"Training images: {train_imgs}, masks: {train_masks}")  # 应均为 800

注意： drive.mount() 后必须手动点击授权链接并粘贴验证码，这是 Colab 的安全机制，无法跳过。若忘记此步，后续所有 !cp 命令将报错 No such file or directory 。

4.2 数据预处理：TFRecord 生成与验证

将原始 PNG 数据转换为 TFRecord 是性能优化的第一步。我们编写专用脚本，关键点在于：

• 并行化处理 ：使用 concurrent.futures.ThreadPoolExecutor 启动 4 个线程，避免 GIL 限制
• 压缩存储 ： tf.train.BytesList 存储 JPEG 压缩后的图像（质量 95），体积减少 62%
• 特征对齐 ：确保 image 和 mask 的 height/width/channels 特征完全一致，避免 tf.io.parse_single_example 解析失败

import tensorflow as tf
import numpy as np
from concurrent.futures import ThreadPoolExecutor
import cv2

def _bytes_feature(value):
    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))

def _int64_feature(value):
    return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))

def create_tfrecord_example(img_path, mask_path):
    # 读取并压缩图像
    img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
    _, img_encoded = cv2.imencode('.jpg', img, [cv2.IMWRITE_JPEG_QUALITY, 95])
    
    mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
    _, mask_encoded = cv2.imencode('.jpg', mask, [cv2.IMWRITE_JPEG_QUALITY, 95])
    
    # 构建 Example
    feature = {
        'image': _bytes_feature(img_encoded.tobytes()),
        'mask': _bytes_feature(mask_encoded.tobytes()),
        'height': _int64_feature(img.shape[0]),
        'width': _int64_feature(img.shape[1]),
        'channels': _int64_feature(1)
    }
    return tf.train.Example(features=tf.train.Features(feature=feature))

# 生成 train.tfrecord
with tf.io.TFRecordWriter('/content/data/train.tfrecord') as writer:
    img_files = sorted(os.listdir('/content/data/covid_ct/train/images/'))
    with ThreadPoolExecutor(max_workers=4) as executor:
        futures = []
        for i, img_file in enumerate(img_files[:100]):  # 先处理 100 张验证流程
            mask_file = img_file.replace('img', 'mask')
            future = executor.submit(create_tfrecord_example,
                                   f'/content/data/covid_ct/train/images/{img_file}',
                                   f'/content/data/covid_ct/train/masks/{mask_file}')
            futures.append(future)
        
        for future in futures:
            example = future.result()
            writer.write(example.SerializeToString())

print("TFRecord generation completed.")

验证 TFRecord 可读性：

# 读取第一条记录并显示
raw_dataset = tf.data.TFRecordDataset('/content/data/train.tfrecord')
for raw_record in raw_dataset.take(1):
    example = tf.train.Example()
    example.ParseFromString(raw_record.numpy())
    print("Height:", example.features.feature['height'].int64_list.value[0])
    print("Image size:", len(example.features.feature['image'].bytes_list.value[0]))

4.3 数据管道构建：从 TFRecord 到 Batched Dataset

构建高效 tf.data.Dataset 是 Colab 性能的生命线。我们采用四级流水线：

def parse_tfrecord(example_proto):
    feature_description = {
        'image': tf.io.FixedLenFeature([], tf.string),
        'mask': tf.io.FixedLenFeature([], tf.string),
        'height': tf.io.FixedLenFeature([], tf.int64),
        'width': tf.io.FixedLenFeature([], tf.int64),
        'channels': tf.io.FixedLenFeature([], tf.int64),
    }
    parsed = tf.io.parse_single_example(example_proto, feature_description)
    
    # 解码 JPEG 并归一化
    image = tf.io.decode_jpeg(parsed['image'], channels=1)
    image = tf.cast(image, tf.float32) / 255.0
    
    mask = tf.io.decode_jpeg(parsed['mask'], channels=1)
    mask = tf.cast(mask, tf.float32) / 255.0
    
    # 裁剪至 512x512（CT 图像通常为 512x512，此步确保尺寸统一）
    image = tf.image.resize(image, [512, 512], method='bilinear')
    mask = tf.image.resize(mask, [512, 512], method='nearest')
    
    return image, mask

def augment_data(image, mask):
    # 随机旋转 ±15°
    angle = tf.random.uniform([], minval=-15, maxval=15, dtype=tf.float32)
    image = tfa.image.rotate(image, angle * np.pi / 180, interpolation='bilinear')
    mask = tfa.image.rotate(mask, angle * np.pi / 180, interpolation='nearest')
    
    # 随机平移
    shift = tf.random.uniform([2], minval=-0.1, maxval=0.1)
    image = tfa.image.translate(image, shift, interpolation='bilinear')
    mask = tfa.image.translate(mask, shift, interpolation='nearest')
    
    return image, mask

# 构建最终 Dataset
dataset = tf.data.TFRecordDataset('/content/data/train.tfrecord')
dataset = dataset.map(parse_tfrecord, num_parallel_calls=tf.data.AUTOTUNE)
dataset = dataset.map(augment_data, num_parallel_calls=tf.data.AUTOTUNE)
dataset = dataset.batch(4)  # batch_size=4 为 T4 最优值
dataset = dataset.prefetch(tf.data.AUTOTUNE)  # 关键！预取下一个 batch

# 验证 pipeline 速度
import time
start = time.time()
for i, (x, y) in enumerate(dataset.take(10)):
    pass
end = time.time()
print(f"10 batches processed in {end-start:.2f}s → {10/(end-start):.1f} batches/sec")

实测结果：在 T4 GPU 上，此 pipeline 达到 8.7 batches/sec，远超原始 Drive 直读的 1.2 batches/sec。

4.4 U-Net 模型构建：Keras Functional API 实现

我们采用 Keras Functional API 实现标准 U-Net，重点优化内存布局：

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers

def conv_block(x, filters, kernel_size=3, activation='relu'):
    x = layers.Conv2D(filters, kernel_size, padding='same', 
                      kernel_initializer='he_normal')(x)
    x = layers.BatchNormalization()(x)
    x = layers.Activation(activation)(x)
    return x

def unet_model(input_shape=(512, 512, 1), num_classes=1):
    inputs = layers.Input(shape=input_shape)
    
    # Encoder path
    c1 = conv_block(inputs, 64)
    c1 = conv_block(c1, 64)
    p1 = layers.MaxPooling2D((2, 2))(c1)  # 256x256
    
    c2 = conv_block(p1, 128)
    c2 = conv_block(c2, 128)
    p2 = layers.MaxPooling2D((2, 2))(c2)  # 128x128
    
    c3 = conv_block(p2, 256)
    c3 = conv_block(c3, 256)
    p3 = layers.MaxPooling2D((2, 2))(c3)  # 64x64
    
    c4 = conv_block(p3, 512)
    c4 = conv_block(c4, 512)
    p4 = layers.MaxPooling2D((2, 2))(c4)  # 32x32
    
    # Bottleneck
    c5 = conv_block(p4, 1024)
    c5 = conv_block(c5, 1024)
    
    # Decoder path
    u6 = layers.Conv2DTranspose(512, (2, 2), strides=(2, 2), padding='same')(c5)
    u6 = layers.concatenate([u6, c4])
    c6 = conv_block(u6, 512)
    c6 = conv_block(c6, 512)
    
    u7 = layers.Conv2DTranspose(256, (2, 2), strides=(2, 2), padding='same')(c6)
    u7 = layers.concatenate([u7, c3])
    c7 = conv_block(u7, 256)
    c7 = conv_block(c7, 256)
    
    u8 = layers.Conv2DTranspose(128, (2, 2), strides=(2, 2), padding='same')(c7)
    u8 = layers.concatenate([u8, c2])
    c8 = conv_block(u8, 128)
    c8 = conv_block(c8, 128)
    
    u9 = layers.Conv2DTranspose(64, (2, 2), strides=(2, 2), padding='same')(c8)
    u9 = layers.concatenate([u9, c1])
    c9 = conv_block(u9, 64)
    c9 = conv_block(c9, 64)
    
    # Output layer
    outputs = layers.Conv2D(num_classes, (1, 1), activation='sigmoid')(c9)
    
    model = keras.Model(inputs=[inputs], outputs=[outputs])
    return model

# 构建模型并打印结构
model = unet_model()
model.summary()

提示： model.summary() 输出的参数量为 31,032,961，显存占用约 9.2GB（T4），符合预期。若 summary 显示 None 形状，说明 input_shape 未正确传递，需检查 layers.Input 的定义。

4.5 模型编译与训练：混合精度与回调配置

启用混合精度训练（FP16）可将 T4 GPU 的训练速度提升 1.8 倍，且不损失精度：

# 启用混合精度
from tensorflow.keras import mixed_precision
policy = mixed_precision.Policy('mixed_float16')
mixed_precision.set_global_policy(policy)

# 重新构建模型（必须在启用 policy 后）
model = unet_model()
model.compile(
    optimizer=keras.optimizers.Adam(learning_rate=3e-4),
    loss=hybrid_loss,  # 使用 3.2 节定义的损失函数
    metrics=['accuracy']
)

# 配置回调
callbacks = [
    # 模型检查点（保存最佳 val_dice）
    keras.callbacks.ModelCheckpoint(
        filepath='/content/drive/MyDrive/models/best_unet.h5',
        monitor='