952 lines
41 KiB
Python
952 lines
41 KiB
Python
import os
|
||
|
||
# XLA multi-threading optimization - MUST be set before importing torch_xla
|
||
# Set these environment variables early to ensure they take effect
|
||
if 'TPU_CORES' in os.environ or 'COLAB_TPU_ADDR' in os.environ:
|
||
# Enable XLA multi-threading for compilation speedup
|
||
os.environ.setdefault('XLA_FLAGS',
|
||
'--xla_cpu_multi_thread_eigen=true ' +
|
||
'--xla_cpu_enable_fast_math=true ' +
|
||
f'--xla_force_host_platform_device_count={os.cpu_count()}'
|
||
)
|
||
# Set PyTorch XLA threading
|
||
os.environ.setdefault('PYTORCH_XLA_COMPILATION_THREADS', str(os.cpu_count()))
|
||
print(f"Set XLA compilation threads to {os.cpu_count()}")
|
||
|
||
import torch
|
||
from torch.utils.data import DataLoader
|
||
from torch.optim.lr_scheduler import LambdaLR
|
||
import random
|
||
import time
|
||
import numpy as np
|
||
import math
|
||
import pathlib
|
||
import logging
|
||
import sys
|
||
import json
|
||
import pickle
|
||
from contextlib import nullcontext
|
||
|
||
from dataset import BrainToTextDataset, train_test_split_indicies
|
||
from data_augmentations import gauss_smooth
|
||
|
||
import torchaudio.functional as F # for edit distance
|
||
from omegaconf import OmegaConf
|
||
|
||
# Import Accelerate for TPU support
|
||
from accelerate import Accelerator, DataLoaderConfiguration
|
||
from accelerate.utils import set_seed
|
||
|
||
# Import XLA after setting environment variables
|
||
import torch_xla.core.xla_model as xm
|
||
|
||
torch.set_float32_matmul_precision('high') # makes float32 matmuls faster on some GPUs
|
||
torch.backends.cudnn.deterministic = True # makes training more reproducible
|
||
torch._dynamo.config.cache_size_limit = 64
|
||
|
||
from rnn_model import TripleGRUDecoder
|
||
|
||
class BrainToTextDecoder_Trainer:
|
||
"""
|
||
This class will initialize and train a brain-to-text phoneme decoder
|
||
|
||
Written by Nick Card and Zachery Fogg with reference to Stanford NPTL's decoding function
|
||
"""
|
||
|
||
def __init__(self, args):
|
||
'''
|
||
args : dictionary of training arguments
|
||
'''
|
||
|
||
# Configure DataLoader behavior for TPU compatibility
|
||
dataloader_config = DataLoaderConfiguration(
|
||
even_batches=False # Required for batch_size=None DataLoaders on TPU
|
||
)
|
||
|
||
# Initialize Accelerator for TPU/multi-device support
|
||
self.use_xla = bool(xm.get_xla_supported_devices())
|
||
self.amp_requested = args.get('use_amp', True)
|
||
mixed_precision_mode = 'bf16' if self.amp_requested else 'no'
|
||
|
||
self.accelerator = Accelerator(
|
||
mixed_precision=mixed_precision_mode,
|
||
gradient_accumulation_steps=args.get('gradient_accumulation_steps', 1),
|
||
log_with=None, # We'll use our own logging
|
||
project_dir=args.get('output_dir', './output'),
|
||
dataloader_config=dataloader_config,
|
||
)
|
||
|
||
|
||
# Trainer fields
|
||
self.args = args
|
||
self.logger = None
|
||
self.device = self.accelerator.device # Use accelerator device instead of manual device selection
|
||
self.model = None
|
||
self.optimizer = None
|
||
self.learning_rate_scheduler = None
|
||
self.ctc_loss = None
|
||
|
||
self.best_val_PER = torch.inf # track best PER for checkpointing
|
||
self.best_val_loss = torch.inf # track best loss for checkpointing
|
||
|
||
self.train_dataset = None
|
||
self.val_dataset = None
|
||
self.train_loader = None
|
||
self.val_loader = None
|
||
|
||
self.transform_args = self.args['dataset']['data_transforms']
|
||
|
||
# Adversarial training config (safe defaults if not provided)
|
||
adv_cfg = self.args.get('adversarial', {})
|
||
self.adv_enabled = adv_cfg.get('enabled', False)
|
||
self.adv_grl_lambda = float(adv_cfg.get('grl_lambda', 0.5)) # GRL strength
|
||
self.adv_noisy_loss_weight = float(adv_cfg.get('noisy_loss_weight', 0.2)) # weight for noisy branch CTC
|
||
self.adv_noise_l2_weight = float(adv_cfg.get('noise_l2_weight', 0.0)) # optional L2 on noise output
|
||
self.adv_warmup_steps = int(adv_cfg.get('warmup_steps', 0)) # delay enabling adversarial after N steps
|
||
|
||
# Create output directory
|
||
if args['mode'] == 'train':
|
||
os.makedirs(self.args['output_dir'], exist_ok=True)
|
||
|
||
# Create checkpoint directory
|
||
if args['save_best_checkpoint'] or args['save_all_val_steps'] or args['save_final_model']:
|
||
os.makedirs(self.args['checkpoint_dir'], exist_ok=True)
|
||
|
||
# Set up logging
|
||
self.logger = logging.getLogger(__name__)
|
||
for handler in self.logger.handlers[:]: # make a copy of the list
|
||
self.logger.removeHandler(handler)
|
||
self.logger.setLevel(logging.INFO)
|
||
formatter = logging.Formatter(fmt='%(asctime)s: %(message)s')
|
||
|
||
if args['mode']=='train':
|
||
# During training, save logs to file in output directory
|
||
fh = logging.FileHandler(str(pathlib.Path(self.args['output_dir'],'training_log')))
|
||
fh.setFormatter(formatter)
|
||
self.logger.addHandler(fh)
|
||
|
||
# Always print logs to stdout
|
||
sh = logging.StreamHandler(sys.stdout)
|
||
sh.setFormatter(formatter)
|
||
self.logger.addHandler(sh)
|
||
|
||
# Log device information (managed by Accelerator)
|
||
self.logger.info(f'Using device: {self.device}')
|
||
self.logger.info(f'Accelerator state: {self.accelerator.state}')
|
||
if self.accelerator.num_processes > 1:
|
||
self.logger.info(f'Distributed training on {self.accelerator.num_processes} processes')
|
||
if self.use_xla and self.amp_requested:
|
||
self.logger.info('AMP requested on TPU; converting model weights to bfloat16 for memory efficiency.')
|
||
|
||
# Set seed if provided (using Accelerator's set_seed for proper distributed seeding)
|
||
if self.args['seed'] != -1:
|
||
set_seed(self.args['seed'])
|
||
|
||
# Initialize the model
|
||
self.model = TripleGRUDecoder(
|
||
neural_dim = self.args['model']['n_input_features'],
|
||
n_units = self.args['model']['n_units'],
|
||
n_days = len(self.args['dataset']['sessions']),
|
||
n_classes = self.args['dataset']['n_classes'],
|
||
rnn_dropout = self.args['model']['rnn_dropout'],
|
||
input_dropout = self.args['model']['input_network']['input_layer_dropout'],
|
||
patch_size = self.args['model']['patch_size'],
|
||
patch_stride = self.args['model']['patch_stride'],
|
||
)
|
||
|
||
if self.use_xla and self.amp_requested:
|
||
self.model = self.model.to(torch.bfloat16)
|
||
self.logger.info('Converted model parameters to bfloat16 for TPU training.')
|
||
|
||
self.model_dtype = next(self.model.parameters()).dtype
|
||
|
||
# Temporarily disable torch.compile for compatibility with new model architecture
|
||
# TODO: Re-enable torch.compile once model is stable
|
||
# self.logger.info("Using torch.compile")
|
||
# self.model = torch.compile(self.model)
|
||
self.logger.info("torch.compile disabled for new TripleGRUDecoder compatibility")
|
||
|
||
self.logger.info(f"Initialized RNN decoding model")
|
||
|
||
self.logger.info(self.model)
|
||
|
||
# Log how many parameters are in the model
|
||
total_params = sum(p.numel() for p in self.model.parameters())
|
||
self.logger.info(f"Model has {total_params:,} parameters")
|
||
|
||
# Determine how many day-specific parameters are in the model
|
||
day_params = 0
|
||
for name, param in self.model.named_parameters():
|
||
if 'day' in name:
|
||
day_params += param.numel()
|
||
|
||
self.logger.info(f"Model has {day_params:,} day-specific parameters | {((day_params / total_params) * 100):.2f}% of total parameters")
|
||
|
||
# Create datasets and dataloaders
|
||
train_file_paths = [os.path.join(self.args["dataset"]["dataset_dir"],s,'data_train.hdf5') for s in self.args['dataset']['sessions']]
|
||
val_file_paths = [os.path.join(self.args["dataset"]["dataset_dir"],s,'data_val.hdf5') for s in self.args['dataset']['sessions']]
|
||
|
||
# Ensure that there are no duplicate days
|
||
if len(set(train_file_paths)) != len(train_file_paths):
|
||
raise ValueError("There are duplicate sessions listed in the train dataset")
|
||
if len(set(val_file_paths)) != len(val_file_paths):
|
||
raise ValueError("There are duplicate sessions listed in the val dataset")
|
||
|
||
# Split trials into train and test sets
|
||
train_trials, _ = train_test_split_indicies(
|
||
file_paths = train_file_paths,
|
||
test_percentage = 0,
|
||
seed = self.args['dataset']['seed'],
|
||
bad_trials_dict = None,
|
||
)
|
||
_, val_trials = train_test_split_indicies(
|
||
file_paths = val_file_paths,
|
||
test_percentage = 1,
|
||
seed = self.args['dataset']['seed'],
|
||
bad_trials_dict = None,
|
||
)
|
||
|
||
# Save dictionaries to output directory to know which trials were train vs val
|
||
with open(os.path.join(self.args['output_dir'], 'train_val_trials.json'), 'w') as f:
|
||
json.dump({'train' : train_trials, 'val': val_trials}, f)
|
||
|
||
# Determine if a only a subset of neural features should be used
|
||
feature_subset = None
|
||
if ('feature_subset' in self.args['dataset']) and self.args['dataset']['feature_subset'] != None:
|
||
feature_subset = self.args['dataset']['feature_subset']
|
||
self.logger.info(f'Using only a subset of features: {feature_subset}')
|
||
|
||
# train dataset and dataloader
|
||
self.train_dataset = BrainToTextDataset(
|
||
trial_indicies = train_trials,
|
||
split = 'train',
|
||
days_per_batch = self.args['dataset']['days_per_batch'],
|
||
n_batches = self.args['num_training_batches'],
|
||
batch_size = self.args['dataset']['batch_size'],
|
||
must_include_days = None,
|
||
random_seed = self.args['dataset']['seed'],
|
||
feature_subset = feature_subset
|
||
)
|
||
# Custom collate function that handles pre-batched data from our dataset
|
||
def collate_fn(batch):
|
||
# Our dataset returns full batches, so batch will be a list of single batch dict
|
||
# Extract the first (and only) element since our dataset.__getitem__() returns a full batch
|
||
if len(batch) == 1 and isinstance(batch[0], dict):
|
||
return batch[0]
|
||
else:
|
||
# Fallback for unexpected batch structure
|
||
return batch
|
||
|
||
# DataLoader configuration compatible with Accelerate
|
||
self.train_loader = DataLoader(
|
||
self.train_dataset,
|
||
batch_size = 1, # Use batch_size=1 since dataset returns full batches
|
||
shuffle = self.args['dataset']['loader_shuffle'],
|
||
num_workers = self.args['dataset']['num_dataloader_workers'],
|
||
pin_memory = True,
|
||
collate_fn = collate_fn
|
||
)
|
||
|
||
# val dataset and dataloader
|
||
self.val_dataset = BrainToTextDataset(
|
||
trial_indicies = val_trials,
|
||
split = 'test',
|
||
days_per_batch = None,
|
||
n_batches = None,
|
||
batch_size = self.args['dataset']['batch_size'],
|
||
must_include_days = None,
|
||
random_seed = self.args['dataset']['seed'],
|
||
feature_subset = feature_subset
|
||
)
|
||
# Validation DataLoader with same collate function
|
||
self.val_loader = DataLoader(
|
||
self.val_dataset,
|
||
batch_size = 1, # Use batch_size=1 since dataset returns full batches
|
||
shuffle = False,
|
||
num_workers = 0, # Keep validation dataloader single-threaded for consistency
|
||
pin_memory = True,
|
||
collate_fn = collate_fn # Use same collate function
|
||
)
|
||
|
||
self.logger.info("Successfully initialized datasets")
|
||
|
||
# Create optimizer, learning rate scheduler, and loss
|
||
self.optimizer = self.create_optimizer()
|
||
|
||
if self.args['lr_scheduler_type'] == 'linear':
|
||
self.learning_rate_scheduler = torch.optim.lr_scheduler.LinearLR(
|
||
optimizer = self.optimizer,
|
||
start_factor = 1.0,
|
||
end_factor = self.args['lr_min'] / self.args['lr_max'],
|
||
total_iters = self.args['lr_decay_steps'],
|
||
)
|
||
elif self.args['lr_scheduler_type'] == 'cosine':
|
||
self.learning_rate_scheduler = self.create_cosine_lr_scheduler(self.optimizer)
|
||
|
||
else:
|
||
raise ValueError(f"Invalid learning rate scheduler type: {self.args['lr_scheduler_type']}")
|
||
|
||
self.ctc_loss = torch.nn.CTCLoss(blank = 0, reduction = 'none', zero_infinity = False)
|
||
|
||
# If a checkpoint is provided, then load from checkpoint
|
||
if self.args['init_from_checkpoint']:
|
||
self.load_model_checkpoint(self.args['init_checkpoint_path'])
|
||
|
||
# Set rnn and/or input layers to not trainable if specified
|
||
for name, param in self.model.named_parameters():
|
||
if not self.args['model']['rnn_trainable'] and 'gru' in name:
|
||
param.requires_grad = False
|
||
|
||
elif not self.args['model']['input_network']['input_trainable'] and 'day' in name:
|
||
param.requires_grad = False
|
||
|
||
# Prepare model, optimizer, scheduler, and dataloaders for distributed training
|
||
# Let Accelerator handle everything automatically for both GPU and TPU
|
||
(
|
||
self.model,
|
||
self.optimizer,
|
||
self.learning_rate_scheduler,
|
||
self.train_loader,
|
||
self.val_loader,
|
||
) = self.accelerator.prepare(
|
||
self.model,
|
||
self.optimizer,
|
||
self.learning_rate_scheduler,
|
||
self.train_loader,
|
||
self.val_loader,
|
||
)
|
||
|
||
self.model_dtype = next(self.model.parameters()).dtype
|
||
|
||
self.logger.info("Prepared model and dataloaders with Accelerator")
|
||
if self.adv_enabled:
|
||
self.logger.info(f"Adversarial training ENABLED | grl_lambda={self.adv_grl_lambda}, noisy_loss_weight={self.adv_noisy_loss_weight}, noise_l2_weight={self.adv_noise_l2_weight}, warmup_steps={self.adv_warmup_steps}")
|
||
|
||
def autocast_context(self):
|
||
"""Return appropriate autocast context; disable on XLA to avoid dtype mismatches."""
|
||
if self.device.type == 'xla':
|
||
return nullcontext()
|
||
return self.accelerator.autocast()
|
||
|
||
def create_optimizer(self):
|
||
'''
|
||
Create the optimizer with special param groups
|
||
|
||
Biases and day weights should not be decayed
|
||
|
||
Day weights should have a separate learning rate
|
||
'''
|
||
bias_params = [p for name, p in self.model.named_parameters() if 'gru.bias' in name or 'out.bias' in name]
|
||
day_params = [p for name, p in self.model.named_parameters() if 'day_' in name]
|
||
other_params = [p for name, p in self.model.named_parameters() if 'day_' not in name and 'gru.bias' not in name and 'out.bias' not in name]
|
||
|
||
if len(day_params) != 0:
|
||
param_groups = [
|
||
{'params' : bias_params, 'weight_decay' : 0, 'group_type' : 'bias'},
|
||
{'params' : day_params, 'lr' : self.args['lr_max_day'], 'weight_decay' : self.args['weight_decay_day'], 'group_type' : 'day_layer'},
|
||
{'params' : other_params, 'group_type' : 'other'}
|
||
]
|
||
else:
|
||
param_groups = [
|
||
{'params' : bias_params, 'weight_decay' : 0, 'group_type' : 'bias'},
|
||
{'params' : other_params, 'group_type' : 'other'}
|
||
]
|
||
|
||
optim = torch.optim.AdamW(
|
||
param_groups,
|
||
lr = self.args['lr_max'],
|
||
betas = (self.args['beta0'], self.args['beta1']),
|
||
eps = self.args['epsilon'],
|
||
weight_decay = self.args['weight_decay'],
|
||
fused = True
|
||
)
|
||
|
||
return optim
|
||
|
||
def create_cosine_lr_scheduler(self, optim):
|
||
lr_max = self.args['lr_max']
|
||
lr_min = self.args['lr_min']
|
||
lr_decay_steps = self.args['lr_decay_steps']
|
||
|
||
lr_max_day = self.args['lr_max_day']
|
||
lr_min_day = self.args['lr_min_day']
|
||
lr_decay_steps_day = self.args['lr_decay_steps_day']
|
||
|
||
lr_warmup_steps = self.args['lr_warmup_steps']
|
||
lr_warmup_steps_day = self.args['lr_warmup_steps_day']
|
||
|
||
def lr_lambda(current_step, min_lr_ratio, decay_steps, warmup_steps):
|
||
'''
|
||
Create lr lambdas for each param group that implement cosine decay
|
||
|
||
Different lr lambda decaying for day params vs rest of the model
|
||
'''
|
||
# Warmup phase
|
||
if current_step < warmup_steps:
|
||
return float(current_step) / float(max(1, warmup_steps))
|
||
|
||
# Cosine decay phase
|
||
if current_step < decay_steps:
|
||
progress = float(current_step - warmup_steps) / float(
|
||
max(1, decay_steps - warmup_steps)
|
||
)
|
||
cosine_decay = 0.5 * (1 + math.cos(math.pi * progress))
|
||
# Scale from 1.0 to min_lr_ratio
|
||
return max(min_lr_ratio, min_lr_ratio + (1 - min_lr_ratio) * cosine_decay)
|
||
|
||
# After cosine decay is complete, maintain min_lr_ratio
|
||
return min_lr_ratio
|
||
|
||
if len(optim.param_groups) == 3:
|
||
lr_lambdas = [
|
||
lambda step: lr_lambda(
|
||
step,
|
||
lr_min / lr_max,
|
||
lr_decay_steps,
|
||
lr_warmup_steps), # biases
|
||
lambda step: lr_lambda(
|
||
step,
|
||
lr_min_day / lr_max_day,
|
||
lr_decay_steps_day,
|
||
lr_warmup_steps_day,
|
||
), # day params
|
||
lambda step: lr_lambda(
|
||
step,
|
||
lr_min / lr_max,
|
||
lr_decay_steps,
|
||
lr_warmup_steps), # rest of model weights
|
||
]
|
||
elif len(optim.param_groups) == 2:
|
||
lr_lambdas = [
|
||
lambda step: lr_lambda(
|
||
step,
|
||
lr_min / lr_max,
|
||
lr_decay_steps,
|
||
lr_warmup_steps), # biases
|
||
lambda step: lr_lambda(
|
||
step,
|
||
lr_min / lr_max,
|
||
lr_decay_steps,
|
||
lr_warmup_steps), # rest of model weights
|
||
]
|
||
else:
|
||
raise ValueError(f"Invalid number of param groups in optimizer: {len(optim.param_groups)}")
|
||
|
||
return LambdaLR(optim, lr_lambdas, -1)
|
||
|
||
def load_model_checkpoint(self, load_path):
|
||
'''
|
||
Load a training checkpoint for distributed training
|
||
'''
|
||
# Load checkpoint on CPU first to avoid OOM issues
|
||
checkpoint = torch.load(load_path, map_location='cpu', weights_only = False) # checkpoint is just a dict
|
||
|
||
# Get unwrapped model for loading state dict
|
||
unwrapped_model = self.accelerator.unwrap_model(self.model)
|
||
unwrapped_model.load_state_dict(checkpoint['model_state_dict'])
|
||
|
||
self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
|
||
self.learning_rate_scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
|
||
self.best_val_PER = checkpoint['val_PER'] # best phoneme error rate
|
||
self.best_val_loss = checkpoint['val_loss'] if 'val_loss' in checkpoint.keys() else torch.inf
|
||
|
||
# Device handling is managed by Accelerator, no need to manually move to device
|
||
|
||
self.logger.info("Loaded model from checkpoint: " + load_path)
|
||
|
||
def save_model_checkpoint(self, save_path, PER, loss):
|
||
'''
|
||
Save a training checkpoint using Accelerator for distributed training
|
||
'''
|
||
# Only save on main process to avoid conflicts
|
||
if self.accelerator.is_main_process:
|
||
# Unwrap model to get base model for saving
|
||
unwrapped_model = self.accelerator.unwrap_model(self.model)
|
||
|
||
checkpoint = {
|
||
'model_state_dict' : unwrapped_model.state_dict(),
|
||
'optimizer_state_dict' : self.optimizer.state_dict(),
|
||
'scheduler_state_dict' : self.learning_rate_scheduler.state_dict(),
|
||
'val_PER' : PER,
|
||
'val_loss' : loss
|
||
}
|
||
|
||
torch.save(checkpoint, save_path)
|
||
|
||
self.logger.info("Saved model to checkpoint: " + save_path)
|
||
|
||
# Save the args file alongside the checkpoint
|
||
with open(os.path.join(self.args['checkpoint_dir'], 'args.yaml'), 'w') as f:
|
||
OmegaConf.save(config=self.args, f=f)
|
||
|
||
# Wait for all processes to complete checkpoint saving
|
||
self.accelerator.wait_for_everyone()
|
||
|
||
def create_attention_mask(self, sequence_lengths):
|
||
|
||
max_length = torch.max(sequence_lengths).item()
|
||
|
||
batch_size = sequence_lengths.size(0)
|
||
|
||
# Create a mask for valid key positions (columns)
|
||
# Shape: [batch_size, max_length]
|
||
key_mask = torch.arange(max_length, device=sequence_lengths.device).expand(batch_size, max_length)
|
||
key_mask = key_mask < sequence_lengths.unsqueeze(1)
|
||
|
||
# Expand key_mask to [batch_size, 1, 1, max_length]
|
||
# This will be broadcast across all query positions
|
||
key_mask = key_mask.unsqueeze(1).unsqueeze(1)
|
||
|
||
# Create the attention mask of shape [batch_size, 1, max_length, max_length]
|
||
# by broadcasting key_mask across all query positions
|
||
attention_mask = key_mask.expand(batch_size, 1, max_length, max_length)
|
||
|
||
# Convert boolean mask to float mask:
|
||
# - True (valid key positions) -> 0.0 (no change to attention scores)
|
||
# - False (padding key positions) -> -inf (will become 0 after softmax)
|
||
attention_mask_float = torch.where(attention_mask,
|
||
True,
|
||
False)
|
||
|
||
return attention_mask_float
|
||
|
||
def transform_data(self, features, n_time_steps, mode = 'train'):
|
||
'''
|
||
Apply various augmentations and smoothing to data
|
||
Performing augmentations is much faster on GPU than CPU
|
||
'''
|
||
|
||
# TPU and GPU should now handle data consistently with our improved DataLoader configuration
|
||
|
||
data_shape = features.shape
|
||
batch_size = data_shape[0]
|
||
channels = data_shape[-1]
|
||
|
||
# We only apply these augmentations in training
|
||
if mode == 'train':
|
||
# add static gain noise
|
||
if self.transform_args['static_gain_std'] > 0:
|
||
warp_mat = torch.tile(torch.unsqueeze(torch.eye(channels), dim = 0), (batch_size, 1, 1))
|
||
warp_mat += torch.randn_like(warp_mat, device=self.device) * self.transform_args['static_gain_std']
|
||
|
||
features = torch.matmul(features, warp_mat)
|
||
|
||
# add white noise
|
||
if self.transform_args['white_noise_std'] > 0:
|
||
features += torch.randn(data_shape, device=self.device) * self.transform_args['white_noise_std']
|
||
|
||
# add constant offset noise
|
||
if self.transform_args['constant_offset_std'] > 0:
|
||
features += torch.randn((batch_size, 1, channels), device=self.device) * self.transform_args['constant_offset_std']
|
||
|
||
# add random walk noise
|
||
if self.transform_args['random_walk_std'] > 0:
|
||
features += torch.cumsum(torch.randn(data_shape, device=self.device) * self.transform_args['random_walk_std'], dim =self.transform_args['random_walk_axis'])
|
||
|
||
# randomly cutoff part of the data timecourse
|
||
if self.transform_args['random_cut'] > 0:
|
||
cut = np.random.randint(0, self.transform_args['random_cut'])
|
||
features = features[:, cut:, :]
|
||
n_time_steps = n_time_steps - cut
|
||
|
||
# Apply Gaussian smoothing to data
|
||
# This is done in both training and validation
|
||
if self.transform_args['smooth_data']:
|
||
features = gauss_smooth(
|
||
inputs = features,
|
||
device = self.device,
|
||
smooth_kernel_std = self.transform_args['smooth_kernel_std'],
|
||
smooth_kernel_size= self.transform_args['smooth_kernel_size'],
|
||
)
|
||
|
||
if hasattr(self, 'model_dtype'):
|
||
features = features.to(self.model_dtype)
|
||
|
||
|
||
return features, n_time_steps
|
||
|
||
def train(self):
|
||
'''
|
||
Train the model
|
||
'''
|
||
|
||
# Set model to train mode (specificially to make sure dropout layers are engaged)
|
||
self.model.train()
|
||
|
||
# create vars to track performance
|
||
train_losses = []
|
||
val_losses = []
|
||
val_PERs = []
|
||
val_results = []
|
||
|
||
val_steps_since_improvement = 0
|
||
|
||
# training params
|
||
save_best_checkpoint = self.args.get('save_best_checkpoint', True)
|
||
early_stopping = self.args.get('early_stopping', True)
|
||
|
||
early_stopping_val_steps = self.args['early_stopping_val_steps']
|
||
|
||
train_start_time = time.time()
|
||
|
||
# train for specified number of batches
|
||
self.logger.info("Starting training loop - loading first batch (TPU compilation may take 5-15 minutes)...")
|
||
for i, batch in enumerate(self.train_loader):
|
||
|
||
self.model.train()
|
||
self.optimizer.zero_grad()
|
||
|
||
# Train step
|
||
start_time = time.time()
|
||
|
||
# Data is automatically moved to device by Accelerator
|
||
features = batch['input_features']
|
||
labels = batch['seq_class_ids']
|
||
n_time_steps = batch['n_time_steps']
|
||
phone_seq_lens = batch['phone_seq_lens']
|
||
day_indicies = batch['day_indicies']
|
||
|
||
# Use Accelerator's autocast (mixed precision handled by Accelerator init)
|
||
with self.autocast_context():
|
||
|
||
# Apply augmentations to the data
|
||
features, n_time_steps = self.transform_data(features, n_time_steps, 'train')
|
||
|
||
# Ensure proper dtype handling for TPU mixed precision
|
||
adjusted_lens = ((n_time_steps.float() - self.args['model']['patch_size']) / self.args['model']['patch_stride'] + 1).to(torch.int32)
|
||
|
||
# Get phoneme predictions using inference mode during training
|
||
# (We use inference mode for simplicity - only clean logits are used for CTC loss)
|
||
# Ensure features tensor matches model parameter dtype for TPU compatibility
|
||
if features.dtype != self.model_dtype:
|
||
features = features.to(self.model_dtype)
|
||
|
||
# Forward pass: enable full adversarial mode if configured and past warmup
|
||
use_full = self.adv_enabled and (i >= self.adv_warmup_steps)
|
||
if use_full:
|
||
clean_logits, noisy_logits, noise_output = self.model(features, day_indicies, None, False, 'full', grl_lambda=self.adv_grl_lambda)
|
||
else:
|
||
logits = self.model(features, day_indicies, None, False, 'inference')
|
||
|
||
# Calculate CTC Loss
|
||
if use_full:
|
||
# Clean CTC loss
|
||
clean_log_probs = torch.permute(clean_logits, [1, 0, 2]).float().log_softmax(2)
|
||
clean_loss = self.ctc_loss(
|
||
clean_log_probs,
|
||
labels,
|
||
adjusted_lens,
|
||
phone_seq_lens
|
||
)
|
||
clean_loss = torch.mean(clean_loss)
|
||
|
||
# Noisy branch CTC loss（让 Noisy 更可识别，但经 GRL 对 NoiseModel 变成对抗）
|
||
noisy_log_probs = torch.permute(noisy_logits, [1, 0, 2]).float().log_softmax(2)
|
||
noisy_loss = self.ctc_loss(
|
||
noisy_log_probs,
|
||
labels,
|
||
adjusted_lens,
|
||
phone_seq_lens
|
||
)
|
||
noisy_loss = torch.mean(noisy_loss)
|
||
|
||
# Optional noise energy regularization
|
||
noise_l2 = torch.tensor(0.0, device=self.device, dtype=clean_loss.dtype)
|
||
if self.adv_noise_l2_weight > 0.0:
|
||
noise_l2 = torch.mean(noise_output.float().pow(2)).to(clean_loss.dtype)
|
||
|
||
loss = clean_loss + self.adv_noisy_loss_weight * noisy_loss + self.adv_noise_l2_weight * noise_l2
|
||
else:
|
||
log_probs = torch.permute(logits, [1, 0, 2]).float().log_softmax(2)
|
||
loss = self.ctc_loss(
|
||
log_probs=log_probs,
|
||
targets=labels,
|
||
input_lengths=adjusted_lens,
|
||
target_lengths=phone_seq_lens
|
||
)
|
||
loss = torch.mean(loss) # take mean loss over batches
|
||
|
||
# Use Accelerator's backward for distributed training
|
||
self.accelerator.backward(loss)
|
||
|
||
# Clip gradient using Accelerator's clip_grad_norm_
|
||
if self.args['grad_norm_clip_value'] > 0:
|
||
grad_norm = self.accelerator.clip_grad_norm_(self.model.parameters(),
|
||
max_norm = self.args['grad_norm_clip_value'])
|
||
|
||
self.optimizer.step()
|
||
self.learning_rate_scheduler.step()
|
||
|
||
# Save training metrics
|
||
train_step_duration = time.time() - start_time
|
||
train_losses.append(loss.detach().item())
|
||
|
||
# Incrementally log training progress
|
||
if i % self.args['batches_per_train_log'] == 0:
|
||
self.logger.info(f'Train batch {i}: ' +
|
||
f'loss: {(loss.detach().item()):.2f} ' +
|
||
f'grad norm: {grad_norm:.2f} '
|
||
f'time: {train_step_duration:.3f}')
|
||
|
||
# Incrementally run a test step
|
||
if i % self.args['batches_per_val_step'] == 0 or i == ((self.args['num_training_batches'] - 1)):
|
||
self.logger.info(f"Running test after training batch: {i}")
|
||
|
||
# Calculate metrics on val data
|
||
start_time = time.time()
|
||
val_metrics = self.validation(loader = self.val_loader, return_logits = self.args['save_val_logits'], return_data = self.args['save_val_data'])
|
||
val_step_duration = time.time() - start_time
|
||
|
||
|
||
# Log info
|
||
self.logger.info(f'Val batch {i}: ' +
|
||
f'PER (avg): {val_metrics["avg_PER"]:.4f} ' +
|
||
f'CTC Loss (avg): {val_metrics["avg_loss"]:.4f} ' +
|
||
f'time: {val_step_duration:.3f}')
|
||
|
||
if self.args['log_individual_day_val_PER']:
|
||
for day in val_metrics['day_PERs'].keys():
|
||
self.logger.info(f"{self.args['dataset']['sessions'][day]} val PER: {val_metrics['day_PERs'][day]['total_edit_distance'] / val_metrics['day_PERs'][day]['total_seq_length']:0.4f}")
|
||
|
||
# Save metrics
|
||
val_PERs.append(val_metrics['avg_PER'])
|
||
val_losses.append(val_metrics['avg_loss'])
|
||
val_results.append(val_metrics)
|
||
|
||
# Determine if new best day. Based on if PER is lower, or in the case of a PER tie, if loss is lower
|
||
new_best = False
|
||
if val_metrics['avg_PER'] < self.best_val_PER:
|
||
self.logger.info(f"New best test PER {self.best_val_PER:.4f} --> {val_metrics['avg_PER']:.4f}")
|
||
self.best_val_PER = val_metrics['avg_PER']
|
||
self.best_val_loss = val_metrics['avg_loss']
|
||
new_best = True
|
||
elif val_metrics['avg_PER'] == self.best_val_PER and (val_metrics['avg_loss'] < self.best_val_loss):
|
||
self.logger.info(f"New best test loss {self.best_val_loss:.4f} --> {val_metrics['avg_loss']:.4f}")
|
||
self.best_val_loss = val_metrics['avg_loss']
|
||
new_best = True
|
||
|
||
if new_best:
|
||
|
||
# Checkpoint if metrics have improved
|
||
if save_best_checkpoint:
|
||
self.logger.info(f"Checkpointing model")
|
||
self.save_model_checkpoint(f'{self.args["checkpoint_dir"]}/best_checkpoint', self.best_val_PER, self.best_val_loss)
|
||
|
||
# save validation metrics to pickle file
|
||
if self.args['save_val_metrics']:
|
||
with open(f'{self.args["checkpoint_dir"]}/val_metrics.pkl', 'wb') as f:
|
||
pickle.dump(val_metrics, f)
|
||
|
||
val_steps_since_improvement = 0
|
||
|
||
else:
|
||
val_steps_since_improvement +=1
|
||
|
||
# Optionally save this validation checkpoint, regardless of performance
|
||
if self.args['save_all_val_steps']:
|
||
self.save_model_checkpoint(f'{self.args["checkpoint_dir"]}/checkpoint_batch_{i}', val_metrics['avg_PER'], val_metrics['avg_loss'])
|
||
|
||
# Early stopping
|
||
if early_stopping and (val_steps_since_improvement >= early_stopping_val_steps):
|
||
self.logger.info(f'Overall validation PER has not improved in {early_stopping_val_steps} validation steps. Stopping training early at batch: {i}')
|
||
break
|
||
|
||
# Log final training steps
|
||
training_duration = time.time() - train_start_time
|
||
|
||
|
||
self.logger.info(f'Best avg val PER achieved: {self.best_val_PER:.5f}')
|
||
self.logger.info(f'Total training time: {(training_duration / 60):.2f} minutes')
|
||
|
||
# Save final model
|
||
if self.args['save_final_model']:
|
||
last_loss = val_losses[-1] if len(val_losses) > 0 else float('inf')
|
||
self.save_model_checkpoint(f'{self.args["checkpoint_dir"]}/final_checkpoint_batch_{i}', val_PERs[-1], last_loss)
|
||
|
||
train_stats = {}
|
||
train_stats['train_losses'] = train_losses
|
||
train_stats['val_losses'] = val_losses
|
||
train_stats['val_PERs'] = val_PERs
|
||
train_stats['val_metrics'] = val_results
|
||
|
||
return train_stats
|
||
|
||
def validation(self, loader, return_logits = False, return_data = False):
|
||
'''
|
||
Calculate metrics on the validation dataset
|
||
'''
|
||
self.model.eval()
|
||
|
||
metrics = {}
|
||
|
||
# Record metrics
|
||
if return_logits:
|
||
metrics['logits'] = []
|
||
metrics['n_time_steps'] = []
|
||
|
||
if return_data:
|
||
metrics['input_features'] = []
|
||
|
||
metrics['decoded_seqs'] = []
|
||
metrics['true_seq'] = []
|
||
metrics['phone_seq_lens'] = []
|
||
metrics['transcription'] = []
|
||
metrics['losses'] = []
|
||
metrics['block_nums'] = []
|
||
metrics['trial_nums'] = []
|
||
metrics['day_indicies'] = []
|
||
|
||
total_edit_distance = 0
|
||
total_seq_length = 0
|
||
|
||
# Calculate PER for each specific day
|
||
day_per = {}
|
||
for d in range(len(self.args['dataset']['sessions'])):
|
||
if self.args['dataset']['dataset_probability_val'][d] == 1:
|
||
day_per[d] = {'total_edit_distance' : 0, 'total_seq_length' : 0}
|
||
|
||
for i, batch in enumerate(loader):
|
||
|
||
# Data is automatically moved to device by Accelerator
|
||
features = batch['input_features']
|
||
labels = batch['seq_class_ids']
|
||
n_time_steps = batch['n_time_steps']
|
||
phone_seq_lens = batch['phone_seq_lens']
|
||
day_indicies = batch['day_indicies']
|
||
|
||
# Determine if we should perform validation on this batch
|
||
day = day_indicies[0].item()
|
||
if self.args['dataset']['dataset_probability_val'][day] == 0:
|
||
if self.args['log_val_skip_logs']:
|
||
self.logger.info(f"Skipping validation on day {day}")
|
||
continue
|
||
|
||
with torch.no_grad():
|
||
|
||
with self.autocast_context():
|
||
features, n_time_steps = self.transform_data(features, n_time_steps, 'val')
|
||
|
||
# Ensure proper dtype handling for TPU mixed precision
|
||
adjusted_lens = ((n_time_steps.float() - self.args['model']['patch_size']) / self.args['model']['patch_stride'] + 1).to(torch.int32)
|
||
|
||
# Ensure features tensor matches model parameter dtype for TPU compatibility
|
||
model_param = next(self.model.parameters()) if self.model is not None else None
|
||
if model_param is not None and features.dtype != model_param.dtype:
|
||
features = features.to(model_param.dtype)
|
||
|
||
logits = self.model(features, day_indicies, None, False, 'inference')
|
||
|
||
val_log_probs = torch.permute(logits, [1, 0, 2]).float().log_softmax(2)
|
||
loss = self.ctc_loss(
|
||
val_log_probs,
|
||
labels,
|
||
adjusted_lens,
|
||
phone_seq_lens,
|
||
)
|
||
loss = torch.mean(loss)
|
||
|
||
metrics['losses'].append(loss.cpu().detach().numpy())
|
||
|
||
# Calculate PER per day and also avg over entire validation set
|
||
batch_edit_distance = 0
|
||
decoded_seqs = []
|
||
for iterIdx in range(logits.shape[0]):
|
||
decoded_seq = torch.argmax(logits[iterIdx, 0 : adjusted_lens[iterIdx], :].clone().detach(),dim=-1)
|
||
decoded_seq = torch.unique_consecutive(decoded_seq, dim=-1)
|
||
decoded_seq = decoded_seq.cpu().detach().numpy()
|
||
decoded_seq = np.array([i for i in decoded_seq if i != 0])
|
||
|
||
trueSeq = np.array(
|
||
labels[iterIdx][0 : phone_seq_lens[iterIdx]].cpu().detach()
|
||
)
|
||
|
||
batch_edit_distance += F.edit_distance(decoded_seq, trueSeq)
|
||
|
||
decoded_seqs.append(decoded_seq)
|
||
|
||
day = batch['day_indicies'][0].item()
|
||
|
||
day_per[day]['total_edit_distance'] += batch_edit_distance
|
||
day_per[day]['total_seq_length'] += torch.sum(phone_seq_lens).item()
|
||
|
||
|
||
total_edit_distance += batch_edit_distance
|
||
total_seq_length += torch.sum(phone_seq_lens)
|
||
|
||
# Record metrics
|
||
if return_logits:
|
||
metrics['logits'].append(logits.cpu().float().numpy()) # Will be in bfloat16 if AMP is enabled, so need to set back to float32
|
||
metrics['n_time_steps'].append(adjusted_lens.cpu().numpy())
|
||
|
||
if return_data:
|
||
metrics['input_features'].append(batch['input_features'].cpu().numpy())
|
||
|
||
metrics['decoded_seqs'].append(decoded_seqs)
|
||
metrics['true_seq'].append(batch['seq_class_ids'].cpu().numpy())
|
||
metrics['phone_seq_lens'].append(batch['phone_seq_lens'].cpu().numpy())
|
||
metrics['transcription'].append(batch['transcriptions'].cpu().numpy())
|
||
metrics['losses'].append(loss.detach().item())
|
||
metrics['block_nums'].append(batch['block_nums'].numpy())
|
||
metrics['trial_nums'].append(batch['trial_nums'].numpy())
|
||
metrics['day_indicies'].append(batch['day_indicies'].cpu().numpy())
|
||
|
||
if isinstance(total_seq_length, torch.Tensor):
|
||
total_length_value = float(total_seq_length.item())
|
||
else:
|
||
total_length_value = float(total_seq_length)
|
||
|
||
avg_PER = total_edit_distance / max(total_length_value, 1e-6)
|
||
|
||
metrics['day_PERs'] = day_per
|
||
metrics['avg_PER'] = avg_PER
|
||
metrics['avg_loss'] = float(np.mean(metrics['losses']))
|
||
|
||
return metrics
|
||
|
||
def inference(self, features, day_indicies, n_time_steps, mode='inference'):
|
||
'''
|
||
TPU-compatible inference method for generating phoneme logits
|
||
'''
|
||
self.model.eval()
|
||
|
||
with torch.no_grad():
|
||
with self.autocast_context():
|
||
# Apply data transformations (no augmentation for inference)
|
||
features, n_time_steps = self.transform_data(features, n_time_steps, 'val')
|
||
|
||
# Ensure features tensor matches model parameter dtype for TPU compatibility
|
||
if features.dtype != self.model_dtype:
|
||
features = features.to(self.model_dtype)
|
||
|
||
# Get phoneme predictions
|
||
logits = self.model(features, day_indicies, None, False, mode)
|
||
|
||
return logits
|
||
|
||
def inference_batch(self, batch, mode='inference'):
|
||
'''
|
||
Inference method for processing a full batch
|
||
'''
|
||
self.model.eval()
|
||
|
||
# Data is automatically moved to device by Accelerator
|
||
features = batch['input_features']
|
||
day_indicies = batch['day_indicies']
|
||
n_time_steps = batch['n_time_steps']
|
||
|
||
with torch.no_grad():
|
||
with self.autocast_context():
|
||
# Apply data transformations (no augmentation for inference)
|
||
features, n_time_steps = self.transform_data(features, n_time_steps, 'val')
|
||
|
||
# Calculate adjusted sequence lengths for CTC with proper dtype handling
|
||
adjusted_lens = ((n_time_steps.float() - self.args['model']['patch_size']) / self.args['model']['patch_stride'] + 1).to(torch.int32)
|
||
|
||
# Ensure features tensor matches model parameter dtype for TPU compatibility
|
||
if features.dtype != self.model_dtype:
|
||
features = features.to(self.model_dtype)
|
||
|
||
# Get phoneme predictions
|
||
logits = self.model(features, day_indicies, None, False, mode)
|
||
|
||
return logits, adjusted_lens |