基于TensorFlow实现PI Model：按样本标签状态动态增减损失函数

Alright, let's walk through implementing the PI Model from the Temporal Ensembling for Semi-Supervised Learning paper using TensorFlow, with the key requirement of dynamically handling loss for labeled and unlabeled samples. The original code was in Theano (now deprecated), so we'll adapt it to modern TensorFlow practices.

Core Concept Recap #

The PI Model's core idea is simple yet effective: feed two augmented versions of the same sample into the exact same neural network, then enforce consistency between the two outputs. For labeled samples, we also add a standard classification loss; for unlabeled samples, only the consistency loss applies.

Step 1: Build the Base Network #

First, we'll replicate the CNN architecture used in the paper (tailored for 32x32 image tasks like CIFAR-10):

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

def build_base_model(input_shape=(32, 32, 3)):
    """Base feature extraction network matching the paper's architecture"""
    return models.Sequential([
        # Block 1
        layers.Conv2D(32, (3,3), activation='relu', padding='same', input_shape=input_shape),
        layers.Conv2D(32, (3,3), activation='relu', padding='same'),
        layers.MaxPooling2D((2,2)),
        layers.Dropout(0.2),
        
        # Block 2
        layers.Conv2D(64, (3,3), activation='relu', padding='same'),
        layers.Conv2D(64, (3,3), activation='relu', padding='same'),
        layers.MaxPooling2D((2,2)),
        layers.Dropout(0.2),
        
        # Block 3
        layers.Conv2D(128, (3,3), activation='relu', padding='same'),
        layers.Conv2D(128, (3,3), activation='relu', padding='same'),
        layers.MaxPooling2D((2,2)),
        layers.Dropout(0.2),
        
        # Classifier head
        layers.Flatten(),
        layers.Dense(512, activation='relu'),
        layers.Dropout(0.5),
        layers.Dense(10, activation='softmax')  # Adjust num classes for your task
    ])

Step 2: Data Augmentation for Consistency #

We need to generate two distinct augmented versions of each sample (as specified in the paper: random flips, crops, and Gaussian noise):

def apply_augmentation(image):
    """Apply random perturbations to a single image"""
    # Random horizontal flip
    image = tf.image.random_flip_left_right(image)
    # Random 2-pixel crop (assumes input is 32x32)
    image = tf.image.random_crop(image, size=[32, 32, 3])
    # Gaussian noise injection
    noise = tf.random.normal(shape=tf.shape(image), mean=0.0, stddev=0.1)
    image = tf.clip_by_value(image + noise, 0.0, 1.0)
    return image

def get_augmented_pair(x):
    """Generate two augmented versions of the same input sample"""
    x1 = apply_augmentation(x)
    x2 = apply_augmentation(x)
    return x1, x2

Step 3: Dynamic Loss Calculation #

This is the critical part: dynamically applying loss based on whether a sample has a label. We'll use a boolean mask (True for labeled samples, False for unlabeled) to filter which samples contribute to the classification loss.

def pi_model_loss(y_true, y_pred1, y_pred2, labeled_mask):
    """
    Calculate combined loss for PI Model:
    - Labeled samples: Cross-entropy + consistency loss
    - Unlabeled samples: Only consistency loss
    
    Args:
        y_true: Ground truth labels (ignored for unlabeled samples)
        y_pred1: Model output for first augmented sample
        y_pred2: Model output for second augmented sample
        labeled_mask: Boolean tensor indicating labeled samples
    """
    # Cross-entropy loss (only for labeled samples)
    ce_loss = tf.keras.losses.sparse_categorical_crossentropy(y_true, y_pred1)
    ce_loss = tf.where(labeled_mask, ce_loss, 0.0)  # Zero out loss for unlabeled samples
    ce_loss = tf.reduce_mean(ce_loss)
    
    # Consistency loss (MSE between the two outputs, applied to all samples)
    consistency_loss = tf.reduce_mean(tf.square(y_pred1 - y_pred2))
    
    # Dynamic lambda weight (linear warmup from 0 to 1, per paper)
    lambda_weight = tf.Variable(0.0, trainable=False)
    total_loss = ce_loss + lambda_weight * consistency_loss
    
    return total_loss, ce_loss, consistency_loss

Step 4: Custom Training Loop #

We'll implement a custom training loop to handle the mixed labeled/unlabeled dataset, update the lambda weight dynamically, and track loss metrics:

def train_pi_model(train_dataset, epochs=100):
    """
    Train the PI Model with dynamic loss handling
    
    Args:
        train_dataset: TensorFlow Dataset yielding (x_batch, y_batch, labeled_mask_batch)
        epochs: Total training epochs
    """
    input_shape = (32, 32, 3)
    model = build_base_model(input_shape)
    optimizer = tf.keras.optimizers.Adam(learning_rate=1e-3)
    
    # Linear warmup for lambda weight (first 10% of epochs, per paper)
    warmup_epochs = int(epochs * 0.1)
    lambda_weight = tf.Variable(0.0, trainable=False)
    
    for epoch in range(epochs):
        # Update lambda weight
        if epoch < warmup_epochs:
            lambda_weight.assign(tf.cast(epoch / warmup_epochs, tf.float32))
        else:
            lambda_weight.assign(1.0)
        
        total_loss_epoch = 0.0
        ce_loss_epoch = 0.0
        consistency_loss_epoch = 0.0
        step_count = 0
        
        for x_batch, y_batch, mask_batch in train_dataset:
            # Generate augmented pairs for the batch
            x1_batch, x2_batch = get_augmented_pair(x_batch)
            
            with tf.GradientTape() as tape:
                # Forward pass both augmented samples through the same model
                y_pred1 = model(x1_batch, training=True)
                y_pred2 = model(x2_batch, training=True)
                
                # Calculate losses
                total_loss, ce_loss, consistency_loss = pi_model_loss(
                    y_batch, y_pred1, y_pred2, mask_batch
                )
            
            # Backward pass and update weights
            gradients = tape.gradient(total_loss, model.trainable_variables)
            optimizer.apply_gradients(zip(gradients, model.trainable_variables))
            
            # Accumulate metrics
            total_loss_epoch += total_loss.numpy()
            ce_loss_epoch += ce_loss.numpy()
            consistency_loss_epoch += consistency_loss.numpy()
            step_count += 1
        
        # Print epoch summary
        print(f"Epoch {epoch+1}/{epochs}")
        print(f"Total Loss: {total_loss_epoch/step_count:.4f} | CE Loss: {ce_loss_epoch/step_count:.4f} | Consistency Loss: {consistency_loss_epoch/step_count:.4f} | Lambda: {lambda_weight.numpy():.4f}\n")
    
    return model

Key Implementation Details #

• Dynamic Loss Filtering: The labeled_mask tensor ensures unlabeled samples don't contribute to cross-entropy loss, which is essential for semi-supervised learning.
• Lambda Warmup: The consistency loss weight starts at 0 and linearly increases to 1 over the first 10% of epochs. This lets the model first learn to fit labeled data before enforcing consistency on unlabeled samples.
• Shared Network: Unlike Temporal Ensembling (which uses a moving average of past predictions), the PI Model uses the same network for both augmented samples, making it simpler to implement and train.

内容的提问来源于stack exchange，提问作者Ola O