train_continue_epoch2.py

#!/usr/bin/env python
"""
Continue training from epoch 2 checkpoint.

This script resumes training from checkpoints/syncnet_fcn_epoch2.pth
which uses SyncNet_TransferLearning with 31-class classification (±15 frames).

Usage:
    python train_continue_epoch2.py --data_dir "E:\voxc2\vox2_dev_mp4_partaa~\dev\mp4" --hours 5
"""

import os
import sys
import argparse
import time
import numpy as np
from pathlib import Path

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
import cv2
import subprocess
from scipy.io import wavfile
import python_speech_features

from SyncNet_TransferLearning import SyncNet_TransferLearning


class AVSyncDataset(Dataset):
    """Dataset for audio-video sync classification."""
    
    def __init__(self, video_dir, max_offset=15, num_samples_per_video=2,
                 frame_size=(112, 112), num_frames=25, max_videos=None):
        self.video_dir = video_dir
        self.max_offset = max_offset
        self.num_samples_per_video = num_samples_per_video
        self.frame_size = frame_size
        self.num_frames = num_frames
        
        # Find all video files
        self.video_files = []
        for ext in ['*.mp4', '*.avi', '*.mov', '*.mkv']:
            self.video_files.extend(Path(video_dir).glob(f'**/{ext}'))
        
        # Limit number of videos if specified
        if max_videos and len(self.video_files) > max_videos:
            np.random.shuffle(self.video_files)
            self.video_files = self.video_files[:max_videos]
        
        if not self.video_files:
            raise ValueError(f"No video files found in {video_dir}")
        
        print(f"Using {len(self.video_files)} video files")
        
        # Generate sample list
        self.samples = []
        for vid_idx in range(len(self.video_files)):
            for _ in range(num_samples_per_video):
                offset = np.random.randint(-max_offset, max_offset + 1)
                self.samples.append((vid_idx, offset))
        
        print(f"Generated {len(self.samples)} training samples")
    
    def __len__(self):
        return len(self.samples)
    
    def extract_features(self, video_path):
        """Extract audio MFCC and video frames."""
        video_path = str(video_path)
        
        # Extract audio
        temp_audio = f'temp_audio_{os.getpid()}_{np.random.randint(10000)}.wav'
        try:
            cmd = ['ffmpeg', '-y', '-i', video_path, '-ac', '1', '-ar', '16000',
                   '-vn', '-acodec', 'pcm_s16le', temp_audio]
            subprocess.run(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, check=True)
            
            sample_rate, audio = wavfile.read(temp_audio)
            
            # Validate audio length
            min_audio_samples = (self.num_frames * 4 + self.max_offset * 4) * 160
            if len(audio) < min_audio_samples:
                raise ValueError(f"Audio too short: {len(audio)} samples")
            
            mfcc = python_speech_features.mfcc(audio, sample_rate, numcep=13)
            
            min_mfcc_frames = self.num_frames * 4 + abs(self.max_offset) * 4
            if len(mfcc) < min_mfcc_frames:
                raise ValueError(f"MFCC too short: {len(mfcc)} frames")
        finally:
            if os.path.exists(temp_audio):
                os.remove(temp_audio)
        
        # Extract video frames
        cap = cv2.VideoCapture(video_path)
        frames = []
        while len(frames) < self.num_frames + abs(self.max_offset) + 10:
            ret, frame = cap.read()
            if not ret:
                break
            frame = cv2.resize(frame, self.frame_size)
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            frames.append(frame.astype(np.float32) / 255.0)
        cap.release()
        
        if len(frames) < self.num_frames + abs(self.max_offset):
            raise ValueError(f"Video too short: {len(frames)} frames")
        
        return mfcc, np.stack(frames)
    
    def apply_offset(self, mfcc, frames, offset):
        """Apply temporal offset between audio and video."""
        mfcc_offset = offset * 4
        
        num_video_frames = min(self.num_frames, len(frames) - abs(offset))
        num_mfcc_frames = num_video_frames * 4
        
        if offset >= 0:
            video_start = 0
            mfcc_start = mfcc_offset
        else:
            video_start = abs(offset)
            mfcc_start = 0
        
        video_segment = frames[video_start:video_start + num_video_frames]
        mfcc_segment = mfcc[mfcc_start:mfcc_start + num_mfcc_frames]
        
        # Pad if needed
        if len(video_segment) < self.num_frames:
            pad_frames = self.num_frames - len(video_segment)
            video_segment = np.concatenate([
                video_segment,
                np.repeat(video_segment[-1:], pad_frames, axis=0)
            ], axis=0)
        
        target_mfcc_len = self.num_frames * 4
        if len(mfcc_segment) < target_mfcc_len:
            pad_mfcc = target_mfcc_len - len(mfcc_segment)
            mfcc_segment = np.concatenate([
                mfcc_segment,
                np.repeat(mfcc_segment[-1:], pad_mfcc, axis=0)
            ], axis=0)
        
        return mfcc_segment[:target_mfcc_len], video_segment[:self.num_frames]
    
    def __getitem__(self, idx):
        vid_idx, offset = self.samples[idx]
        video_path = self.video_files[vid_idx]
        
        try:
            mfcc, frames = self.extract_features(video_path)
            mfcc, frames = self.apply_offset(mfcc, frames, offset)
            
            audio_tensor = torch.FloatTensor(mfcc.T).unsqueeze(0)  # [1, 13, T]
            video_tensor = torch.FloatTensor(frames).permute(3, 0, 1, 2)  # [3, T, H, W]
            offset_tensor = torch.tensor(offset, dtype=torch.long)
            
            return audio_tensor, video_tensor, offset_tensor
        except Exception as e:
            return None


def collate_fn_skip_none(batch):
    """Skip None samples."""
    batch = [b for b in batch if b is not None]
    if len(batch) == 0:
        return None
    
    audio = torch.stack([b[0] for b in batch])
    video = torch.stack([b[1] for b in batch])
    offset = torch.stack([b[2] for b in batch])
    return audio, video, offset


def train_epoch(model, dataloader, criterion, optimizer, device, max_offset):
    """Train for one epoch."""
    model.train()
    total_loss = 0
    total_correct = 0
    total_samples = 0
    
    for batch_idx, batch in enumerate(dataloader):
        if batch is None:
            continue
        
        audio, video, target_offset = batch
        audio = audio.to(device)
        video = video.to(device)
        target_class = (target_offset + max_offset).long().to(device)
        
        optimizer.zero_grad()
        
        # Forward pass
        sync_probs, _, _ = model(audio, video)
        
        # Global average pooling over time
        sync_logits = torch.log(sync_probs + 1e-8).mean(dim=2)  # [B, 31]
        
        # Compute loss
        loss = criterion(sync_logits, target_class)
        
        # Backward pass
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        optimizer.step()
        
        # Track metrics
        total_loss += loss.item() * audio.size(0)
        predicted_class = sync_logits.argmax(dim=1)
        total_correct += (predicted_class == target_class).sum().item()
        total_samples += audio.size(0)
        
        if batch_idx % 10 == 0:
            acc = 100.0 * total_correct / total_samples if total_samples > 0 else 0
            print(f"  Batch {batch_idx}/{len(dataloader)}: Loss={loss.item():.4f}, Acc={acc:.2f}%")
    
    return total_loss / total_samples, total_correct / total_samples


def main():
    parser = argparse.ArgumentParser(description='Continue training from epoch 2')
    parser.add_argument('--data_dir', type=str, required=True)
    parser.add_argument('--checkpoint', type=str, default='checkpoints/syncnet_fcn_epoch2.pth')
    parser.add_argument('--output_dir', type=str, default='checkpoints')
    parser.add_argument('--hours', type=float, default=5.0, help='Training time in hours')
    parser.add_argument('--batch_size', type=int, default=32)
    parser.add_argument('--lr', type=float, default=1e-4)
    parser.add_argument('--max_videos', type=int, default=None, 
                       help='Limit number of videos (for faster training)')
    
    args = parser.parse_args()
    
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"Using device: {device}")
    
    max_offset = 15  # ±15 frames, 31 classes
    
    # Create model
    print("Creating model...")
    model = SyncNet_TransferLearning(
        video_backbone='fcn',
        audio_backbone='fcn', 
        embedding_dim=512,
        max_offset=max_offset,
        freeze_backbone=False
    )
    
    # Load checkpoint
    print(f"Loading checkpoint: {args.checkpoint}")
    checkpoint = torch.load(args.checkpoint, map_location=device)
    
    # Load model state
    model_state = checkpoint['model_state_dict']
    # Remove 'fcn_model.' prefix if present
    new_state = {}
    for k, v in model_state.items():
        if k.startswith('fcn_model.'):
            new_state[k[10:]] = v  # Remove 'fcn_model.' prefix
        else:
            new_state[k] = v
    
    model.load_state_dict(new_state, strict=False)
    start_epoch = checkpoint.get('epoch', 2)
    print(f"Resuming from epoch {start_epoch}")
    
    model = model.to(device)
    
    # Dataset
    print("Loading dataset...")
    dataset = AVSyncDataset(
        video_dir=args.data_dir,
        max_offset=max_offset,
        num_samples_per_video=2,
        max_videos=args.max_videos
    )
    
    dataloader = DataLoader(
        dataset,
        batch_size=args.batch_size,
        shuffle=True,
        num_workers=0,
        collate_fn=collate_fn_skip_none,
        pin_memory=True
    )
    
    # Loss and optimizer
    criterion = nn.CrossEntropyLoss(label_smoothing=0.1)
    optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=1e-4)
    
    # Training loop with time limit
    os.makedirs(args.output_dir, exist_ok=True)
    
    max_seconds = args.hours * 3600
    start_time = time.time()
    epoch = start_epoch
    best_acc = 0
    
    print(f"\n{'='*60}")
    print(f"Starting training for {args.hours} hours...")
    print(f"{'='*60}")
    
    while True:
        elapsed = time.time() - start_time
        remaining = max_seconds - elapsed
        
        if remaining <= 0:
            print(f"\nTime limit reached ({args.hours} hours)")
            break
        
        epoch += 1
        print(f"\nEpoch {epoch} (Time remaining: {remaining/3600:.2f} hours)")
        print("-" * 40)
        
        train_loss, train_acc = train_epoch(
            model, dataloader, criterion, optimizer, device, max_offset
        )
        
        print(f"Epoch {epoch}: Loss={train_loss:.4f}, Acc={100*train_acc:.2f}%")
        
        # Save checkpoint
        checkpoint_path = os.path.join(args.output_dir, f'syncnet_fcn_epoch{epoch}.pth')
        torch.save({
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'loss': train_loss,
            'accuracy': train_acc * 100,
        }, checkpoint_path)
        print(f"Saved: {checkpoint_path}")
        
        # Save best
        if train_acc > best_acc:
            best_acc = train_acc
            best_path = os.path.join(args.output_dir, 'syncnet_fcn_best.pth')
            torch.save({
                'epoch': epoch,
                'model_state_dict': model.state_dict(),
                'optimizer_state_dict': optimizer.state_dict(),
                'loss': train_loss,
                'accuracy': train_acc * 100,
            }, best_path)
            print(f"New best model saved: {best_path}")
    
    print(f"\n{'='*60}")
    print(f"Training complete!")
    print(f"Final epoch: {epoch}")
    print(f"Best accuracy: {100*best_acc:.2f}%")
    print(f"{'='*60}")


if __name__ == '__main__':
    main()