train.py

#!/usr/bin/env python
# coding: utf-8

# Script originally based on
# https://www.kaggle.com/code/philculliton/inferring-birds-with-kaggle-models
#
# Also:
# https://sagemaker.readthedocs.io/en/stable/frameworks/tensorflow/using_tf.html
# 

import argparse
import os
import pickle
import zipfile

import boto3
import keras
import numpy as np
from sagemaker_core.helper.session_helper import Session
import tensorflow as tf
from keras.src.utils import io_utils

import preprocessing
import constants
from birddata import data
from birdmodeling import spectrogram_model
from calibrate_export_model import do_calibrate, save_competition_classes

batch_size = 98*2
num_parallel = 6
monitor_metric = 'val_multi_category_accuracy'


class EpochModelCheckpoint(keras.callbacks.ModelCheckpoint):
    def __init__(self, epoch_file=None, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.epoch_file = epoch_file

    def on_epoch_end(self, epoch, logs=None):
        before_mtime = self._checkpoint_mtime()
        super().on_epoch_end(epoch, logs)
        after_mtime = self._checkpoint_mtime()
        if self.epoch_file:
            new_checkpoint = (after_mtime != before_mtime)
            if new_checkpoint:
                pickle.dump(epoch+1, open(self.epoch_file, 'wb'))

    def _checkpoint_mtime(self):
        if os.path.exists(self.filepath):
            return os.stat(self.filepath).st_mtime
        else:
            return None

class EarlyStoppingPrintBest(tf.keras.callbacks.EarlyStopping):
    def on_train_end(self, logs=None):
        super().on_train_end(logs)
        if self.verbose > 0:
            io_utils.print_msg(
                f"Restored model weights have {self.monitor}: {self.best}"
            )

def get_parser():
    parser = argparse.ArgumentParser()
    parser.add_argument('--learning_rate', type=float)
    parser.add_argument('--time_mask_param', type=int)
    parser.add_argument('--freq_mask_param', type=int)
    parser.add_argument('--reduce_db_param', type=float)
    parser.add_argument('--feat_drop_rate', type=float)
    parser.add_argument('--pos_weight', type=float)
    parser.add_argument('--model_dir', type=str, default=os.environ.get('SM_MODEL_DIR'))
    parser.add_argument('--train', type=str, default=os.environ.get('SM_CHANNEL_TRAIN'))
    parser.add_argument('--checkpoint_path', type=str, default='/opt/ml/checkpoints')
    parser.add_argument('--epochs', type=int, default=200)
    return parser


def get_data(data_dir, zip_file, csv_file, extract_dir='/tmp/data'):
    local_zip = os.path.join(data_dir, zip_file)

    if not os.path.exists(local_zip):
        boto = boto3.Session(region_name=os.environ['AWS_DEFAULT_REGION'])
        sagemaker_session = Session(boto_session=boto)
        bucket = sagemaker_session.default_bucket()
        boto.client('s3').download_file(bucket, zip_file, local_zip)
    zipfile.ZipFile(local_zip).extractall(extract_dir)

    metadata = data.Data(os.path.join(extract_dir, csv_file))
    fold_dict = metadata.get_folds()
    # files_for_combining = data.get_train_filenames_per_index(fold_dict['train'], metadata.competition_classes)
    filename_datasets = {k: tf.data.Dataset.from_tensor_slices((fold_dict[k].full_path,
                                                                np.stack(fold_dict[k].all_indexes)))
                         for k in ['train', 'valid']}

    class_weight = fold_dict['class_weight']

    audio_datasets_batch_raw = {
        'train': filename_datasets['train'].shuffle(len(filename_datasets['train']), reshuffle_each_iteration=True)
                                           .batch(batch_size)
                                           .map(lambda x, y: (preprocessing.load_audio_vectorize(x, fixed_length=constants.fixed_num_samples, augment=True), y), num_parallel_calls=num_parallel)
        ,
        'valid': filename_datasets['valid'].batch(32).map(lambda x, y: (preprocessing.load_audio_vectorize(x, fixed_length=None), y), num_parallel_calls=num_parallel)
    }
    audio_datasets_batch = {
        'train': audio_datasets_batch_raw['train'].map(lambda x, y: preprocessing.combine_recordings_within_batch(x, y), num_parallel_calls=num_parallel)
                                                  #.map(lambda x, y: (preprocessing.combine_recordings_vectorize(x, y, files_for_combining), y), num_parallel_calls=num_parallel)
                                                  .map(lambda x, y: (preprocessing.frame_audio_vectorize(x), y), num_parallel_calls=num_parallel)
        ,
        'valid': audio_datasets_batch_raw['valid'].map(lambda x, y: (preprocessing.frame_audio_vectorize(x), y), num_parallel_calls=num_parallel),
    }
    return {'metadata': metadata,
            'fold_dict': fold_dict,
            'class_weight': class_weight,
            'audio_datasets_batch': audio_datasets_batch,
            'filename_datasets':  filename_datasets}


def train(args):

    print("ARGS", args)
    train_data = get_data(args.train,
                          'birdclef-2021.zip',
                          'train_metadata.csv')
    train_batches = train_data['audio_datasets_batch']['train']
    valid_batches = train_data['audio_datasets_batch']['valid']
    class_weight = train_data['class_weight']
    num_classes = train_data['metadata'].num_classes

    stop_early_peak = EarlyStoppingPrintBest(
        monitor=monitor_metric,
        mode='max',
        min_delta=0.01,
        patience=5,
        restore_best_weights=True,
        verbose=1,
    )
    plateau = tf.keras.callbacks.ReduceLROnPlateau(
        monitor=monitor_metric,
        factor=.5,
        patience=2,
        verbose=1,
        mode='max',
        min_delta=0.01,
        cooldown=1,
    )

    static_checkpoint_file = os.path.join(args.checkpoint_path, 'checkpoint.weights.h5')
    static_checkpoint_epoch_file = os.path.join(args.checkpoint_path, 'checkpoint_epoch.pickle')

    # static_checkpoint = keras.callbacks.ModelCheckpoint(filepath=static_checkpoint_file, save_weights_only=True, verbose=1)
    static_checkpoint = EpochModelCheckpoint(epoch_file=static_checkpoint_epoch_file,
                                             filepath=static_checkpoint_file,
                                             monitor=monitor_metric,
                                             save_best_only=True,
                                             save_weights_only=True,
                                             mode='max')

    spec_model = spectrogram_model.setup_model(
        num_classes,
        weights='imagenet',
        learning_rate=args.learning_rate,
        time_mask_param=args.time_mask_param,
        freq_mask_param=args.freq_mask_param,
        reduce_db_param=args.reduce_db_param,
        feat_drop_rate=args.feat_drop_rate,
        pos_weight=args.pos_weight,
        trainable_start=None,
        gauss_noise_param=None,
        freq_drop_rate=None,
        do_freq_fade=False,
        do_time_fade=False,
    )

    if os.path.exists(static_checkpoint_file):
        initial_epoch = pickle.load(open(static_checkpoint_epoch_file, 'rb'))
        print(f"Loading checkpoint for epoch {initial_epoch}:", static_checkpoint_file)
        spec_model.load_weights(static_checkpoint_file)
    else:
        initial_epoch = 0

    fit_history = spec_model.fit(
        train_batches.prefetch(1),
        validation_data=valid_batches.prefetch(1),
        verbose=2,
        epochs=args.epochs,
        callbacks=[stop_early_peak, plateau, static_checkpoint],
        class_weight=class_weight,
        initial_epoch=initial_epoch,
    )

    # https://github.com/aws/sagemaker-python-sdk/issues/599
    spec_model.export(os.path.join(args.model_dir, 'export/Servo/1'), verbose=False)
    cl = do_calibrate(spec_model, train_data['fold_dict']['valid'], num_classes)
    calib_file = os.path.join(args.model_dir, 'export_calib.pkl')
    classes_file = os.path.join(args.model_dir, 'competition_classes.txt')
    pickle.dump(cl, open(calib_file, 'wb'))
    save_competition_classes(classes_file, train_data['metadata'])
    return fit_history


if __name__ == "__main__":
    parse = get_parser()
    args, _ = parse.parse_known_args()
    train(args)