{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":21154,"databundleVersionId":1243559,"sourceType":"competition"}],"dockerImageVersionId":31259,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import tensorflow as tf\nimport numpy as np\nimport os, re\nimport matplotlib.pyplot as plt\nfrom tensorflow.keras import mixed_precision\n\n# -----------------------------------------------------------\n# 1. 硬件配置：自动识别 TPU 或 GPU\n# -----------------------------------------------------------\ntry:\n    # 尝试检测 TPU\n    tpu = tf.distribute.cluster_resolver.TPUClusterResolver()\n    tf.config.experimental_connect_to_cluster(tpu)\n    tf.tpu.experimental.initialize_tpu_system(tpu)\n    strategy = tf.distribute.TPUStrategy(tpu)\n    print(\"Running on TPU\")\nexcept ValueError:\n    # 如果没有 TPU，检测 GPU\n    gpus = tf.config.list_physical_devices('GPU')\n    if gpus:\n        strategy = tf.distribute.MirroredStrategy()\n        print(f\"Running on GPU. Replicas: {strategy.num_replicas_in_sync}\")\n    else:\n        strategy = tf.distribute.get_strategy()\n        print(\"Running on CPU\")\n\n# -----------------------------------------------------------\n# 2. 开启混合精度训练 (针对 GPU 加速，TPU 默认支持 bfloat16)\n# -----------------------------------------------------------\nif isinstance(strategy, tf.distribute.MirroredStrategy):\n    policy = mixed_precision.Policy('mixed_float16')\n    mixed_precision.set_global_policy(policy)\n    print('Mixed precision enabled for GPU')\n\n# -----------------------------------------------------------\n# 3. 路径与参数设置 (Kaggle 专用路径)\n# -----------------------------------------------------------\nDATA_PATH = '/kaggle/input/tpu-getting-started/tfrecords-jpeg-224x224'\n\nIMAGE_SIZE = [224, 224] \nEPOCHS = 25 # 专家建议：全微调模型可以适当增加 Epoch \nBATCH_SIZE = 16 * strategy.num_replicas_in_sync\n\nTRAINING_FILENAMES = tf.io.gfile.glob(DATA_PATH + '/train/*.tfrec')\nVALIDATION_FILENAMES = tf.io.gfile.glob(DATA_PATH + '/val/*.tfrec')\nTEST_FILENAMES = tf.io.gfile.glob(DATA_PATH + '/test/*.tfrec')\n\nNUM_CLASSES = 104\n\n# -----------------------------------------------------------\n# 4. 数据管道 (保持高精度的核心：数据增强)\n# -----------------------------------------------------------\ndef decode_image(image_data):\n    image = tf.image.decode_jpeg(image_data, channels=3)\n    image = tf.cast(image, tf.float32) / 255.0  \n    image = tf.reshape(image, [*IMAGE_SIZE, 3])\n    return image\n\ndef read_labeled_tfrecord(example):\n    LABELED_TFREC_FORMAT = {\n        \"image\": tf.io.FixedLenFeature([], tf.string),\n        \"class\": tf.io.FixedLenFeature([], tf.int64),\n    }\n    example = tf.io.parse_single_example(example, LABELED_TFREC_FORMAT)\n    image = decode_image(example['image'])\n    label = tf.cast(example['class'], tf.int32)\n    label = tf.one_hot(label, NUM_CLASSES)\n    return image, label\n\ndef read_unlabeled_tfrecord(example):\n    UNLABELED_TFREC_FORMAT = {\n        \"image\": tf.io.FixedLenFeature([], tf.string),\n        \"id\": tf.io.FixedLenFeature([], tf.string),\n    }\n    example = tf.io.parse_single_example(example, UNLABELED_TFREC_FORMAT)\n    image = decode_image(example['image'])\n    idnum = example['id']\n    return image, idnum\n\ndef data_augment(image, label):\n    # 高精度技巧：引入更多随机变换\n    image = tf.image.random_flip_left_right(image)\n    image = tf.image.random_saturation(image, 0.7, 1.3)\n    image = tf.image.random_contrast(image, 0.8, 1.2)\n    image = tf.image.random_brightness(image, 0.1)\n    return image, label\n\ndef get_training_dataset():\n    dataset = tf.data.TFRecordDataset(TRAINING_FILENAMES, num_parallel_reads=tf.data.AUTOTUNE)\n    dataset = dataset.map(read_labeled_tfrecord, num_parallel_calls=tf.data.AUTOTUNE)\n    dataset = dataset.map(data_augment, num_parallel_calls=tf.data.AUTOTUNE)\n    dataset = dataset.repeat().shuffle(2048).batch(BATCH_SIZE).prefetch(tf.data.AUTOTUNE)\n    return dataset\n\ndef get_validation_dataset(ordered=False):\n    dataset = tf.data.TFRecordDataset(VALIDATION_FILENAMES).map(read_labeled_tfrecord).batch(BATCH_SIZE).cache().prefetch(tf.data.AUTOTUNE)\n    return dataset\n\n# -----------------------------------------------------------\n# 5. 模型构建与学习率调度器\n# -----------------------------------------------------------\ndef lr_fn(epoch):\n    LR_START = 0.00001\n    LR_MAX = 0.00005 * strategy.num_replicas_in_sync\n    LR_MIN = 0.00001\n    LR_RAMPUP_EPOCHS = 5\n    LR_SUSTAIN_EPOCHS = 0\n    LR_EXP_DECAY = .8\n    if epoch < LR_RAMPUP_EPOCHS:\n        lr = (LR_MAX - LR_START) / LR_RAMPUP_EPOCHS * epoch + LR_START\n    elif epoch < LR_RAMPUP_EPOCHS + LR_SUSTAIN_EPOCHS:\n        lr = LR_MAX\n    else:\n        lr = (LR_MAX - LR_MIN) * LR_EXP_DECAY**(epoch - LR_RAMPUP_EPOCHS - LR_SUSTAIN_EPOCHS) + LR_MIN\n    return lr\n\nlr_callback = tf.keras.callbacks.LearningRateScheduler(lr_fn, verbose=True)\n\nwith strategy.scope():    \n    pretrained_model = tf.keras.applications.DenseNet201(\n        weights='imagenet', \n        include_top=False, \n        input_shape=[*IMAGE_SIZE, 3]\n    )\n    pretrained_model.trainable = True \n    \n    model = tf.keras.Sequential([\n        pretrained_model,\n        tf.keras.layers.GlobalAveragePooling2D(),\n        tf.keras.layers.Dropout(0.3),\n        # 强制输出层使用 float32 保证 Softmax 的数值稳定性\n        tf.keras.layers.Dense(NUM_CLASSES, activation='softmax', dtype='float32')\n    ])\n        \n    model.compile(\n        optimizer=tf.keras.optimizers.Adam(), \n        loss = tf.keras.losses.CategoricalCrossentropy(label_smoothing=0.1),\n        metrics=['accuracy']\n    )\n\n# -----------------------------------------------------------\n# 6. 执行训练\n# -----------------------------------------------------------\ndef count_data_items(filenames):\n    n = [int(re.compile(r\"-([0-9]+)\\.\").search(fname).group(1)) for fname in filenames]\n    return np.sum(n)\n\nNUM_TRAINING_IMAGES = count_data_items(TRAINING_FILENAMES)\nprint(f\"Training images: {NUM_TRAINING_IMAGES}\")\n\nhistory = model.fit(\n    get_training_dataset(), \n    steps_per_epoch=NUM_TRAINING_IMAGES // BATCH_SIZE, \n    epochs=EPOCHS, \n    validation_data=get_validation_dataset(),\n    callbacks=[lr_callback]\n)\n\n# -----------------------------------------------------------\n# 7. 预测逻辑 (注意：Kaggle 环境需要处理 decode 字节串问题)\n# -----------------------------------------------------------\ntest_ds = tf.data.TFRecordDataset(TEST_FILENAMES).map(read_unlabeled_tfrecord).batch(BATCH_SIZE)\ntest_images_ds = test_ds.map(lambda image, idnum: image)\nprobabilities = model.predict(test_images_ds)\npredictions = np.argmax(probabilities, axis=-1)\n\nprint('Generating submission.csv...')\ntest_ids_ds = test_ds.map(lambda image, idnum: idnum).unbatch()\ntest_ids = next(iter(test_ids_ds.batch(count_data_items(TEST_FILENAMES)))).numpy().astype('U') \n\nimport pandas as pd\nsubmission = pd.DataFrame({'id': test_ids, 'label': predictions})\nsubmission.to_csv('submission.csv', index=False)\nprint(\"Submission saved!\")","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null}]}