Bootcamp

# ===== Objectives =====
# 1. Load a small image dataset
# 2. Build a model (transfer learning)
# 3. Train with callbacks + mixed precision
# 4. Evaluate
# 5. Export to SavedModel + TFLite

# ===== 0-5 min: setup =====
# pip install tensorflow tensorflow-datasets matplotlib
import tensorflow as tf
import tensorflow_datasets as tfds

# Mixed precision if you have a modern GPU
# tf.keras.mixed_precision.set_global_policy('mixed_float16')

# ===== 5-15 min: dataset =====
(ds_train, ds_val), info = tfds.load(
    'tf_flowers', split=['train[:90%]', 'train[90%:]'],
    as_supervised=True, with_info=True,
)

IMG = 224
BATCH = 32
NUM_CLASSES = info.features['label'].num_classes

def prep(image, label):
    image = tf.image.resize(image, (IMG, IMG))
    image = tf.cast(image, tf.float32)
    return image, label

ds_train = (ds_train.map(prep).cache().shuffle(1000).batch(BATCH).prefetch(tf.data.AUTOTUNE))
ds_val   = (ds_val.map(prep).batch(BATCH).prefetch(tf.data.AUTOTUNE))

# ===== 15-25 min: model =====
augment = tf.keras.Sequential([
    tf.keras.layers.RandomFlip('horizontal'),
    tf.keras.layers.RandomRotation(0.1),
    tf.keras.layers.RandomZoom(0.1),
])

base = tf.keras.applications.MobileNetV3Small(
    input_shape=(IMG, IMG, 3), include_top=False, weights='imagenet',
)
base.trainable = False

inputs = tf.keras.Input(shape=(IMG, IMG, 3))
x = augment(inputs)
x = tf.keras.applications.mobilenet_v3.preprocess_input(x)
x = base(x, training=False)
x = tf.keras.layers.GlobalAveragePooling2D()(x)
x = tf.keras.layers.Dropout(0.2)(x)
outputs = tf.keras.layers.Dense(NUM_CLASSES, dtype='float32')(x)
model = tf.keras.Model(inputs, outputs)

model.compile(
    optimizer=tf.keras.optimizers.AdamW(1e-3),
    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    metrics=['accuracy'],
)

# ===== 25-40 min: train =====
cbs = [
    tf.keras.callbacks.EarlyStopping(monitor='val_loss', patience=3, restore_best_weights=True),
    tf.keras.callbacks.ModelCheckpoint('best.keras', save_best_only=True),
    tf.keras.callbacks.ReduceLROnPlateau(monitor='val_loss', patience=2, factor=0.5),
]

history = model.fit(ds_train, validation_data=ds_val, epochs=10, callbacks=cbs)

# ===== 40-50 min: fine-tune top layers of base =====
base.trainable = True
for layer in base.layers[:-30]:
    layer.trainable = False

model.compile(
    optimizer=tf.keras.optimizers.AdamW(1e-5),
    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    metrics=['accuracy'],
)
model.fit(ds_train, validation_data=ds_val, epochs=5, callbacks=cbs)

# ===== 50-55 min: evaluate =====
loss, acc = model.evaluate(ds_val)
print(f'val acc: {acc:.4f}')

# ===== 55-60 min: export =====
# SavedModel for TF Serving / TF Hub / production
model.export('saved_model_dir')

# TFLite for mobile + edge
conv = tf.lite.TFLiteConverter.from_keras_model(model)
conv.optimizations = [tf.lite.Optimize.DEFAULT]
tflite_model = conv.convert()
open('flowers.tflite', 'wb').write(tflite_model)

# Quick inference via TFLite Interpreter
itp = tf.lite.Interpreter(model_content=tflite_model)
itp.allocate_tensors()
print('input details:', itp.get_input_details()[0]['shape'])
print('output classes:', NUM_CLASSES)

# ===== Post-bootcamp checklist =====
# - Train acc + val acc reported on a held-out split
# - Best checkpoint saved (not just the last epoch)
# - Mixed precision on if GPU supports it
# - SavedModel exported (for serving)
# - TFLite model exported (for mobile)
# - README that says how to retrain in one command

# ===== Pitfalls =====
# - softmax + from_logits=True together -> training degrades
# - missing prefetch -> GPU idle
# - augment INSIDE the model means it runs at inference time
#   (you usually want it ONLY during training; use training=True logic)
# - One-shot save with .h5 instead of .keras format

# 30-day TF bootcamp in the lesson body.