training.py

"""
Copyright (c) 2025 Samsung Electronics Co., Ltd.

Author(s):
Mahmoud Afifi (m.afifi1@samsung.com, m.3afifi@gmail.com)

Licensed under the Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0) License, (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at https://creativecommons.org/licenses/by-nc/4.0
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and limitations under the License.
For conditions of distribution and use, see the accompanying LICENSE.md file.

Training script.
"""

import torch.optim as optim
from torch.utils.data import DataLoader
import argparse

import utils
from model import IllumEstimator
from dataloader import IlluminantEstimationDataLoader
from torch.optim.lr_scheduler import LambdaLR, CosineAnnealingLR
from utils import *

torch.manual_seed(RANDOM_SEED)
torch.backends.cudnn.deterministic = True
np.random.seed(RANDOM_SEED)

def get_args():
  parser = argparse.ArgumentParser(
    description='Script for training a model.'
  )
  parser.add_argument('--dataset_path', type=str, required=True,
                      help='Path to the dataset directory. The directory must contain subdirectories:'
                           ' "train" and "val".')
  parser.add_argument('--output_path', type=str, default='./models',
                      help='Path to the output directory for saving the final trained model.')
  parser.add_argument('--capture_data', type=str, nargs='+', required=True,
                      choices=CAPTURE_DATA,
                      help=(
                        'List of capture data features to be used by the model (e.g., --capture_data iso flash). '
                        'Options: "iso", "shutter_speed", "flash", "time", "noise_stats".'))
  parser.add_argument('--target_size', type=int, nargs='+', default=None,
                      help='Target size of raw images, specified as two integers: height width.')
  parser.add_argument('--learning_rate', type=float, default=1e-3,
                      help='Learning rate for the optimizer.')
  parser.add_argument('--weight_decay', type=float, default=1e-9,
                      help='L2 weight decay factor.')
  parser.add_argument('--batch_size', type=int, default=8,
                      help='Mini-batch size.')
  parser.add_argument('--epochs', type=int, default=400,
                      help='Number of epochs to train the model.')
  parser.add_argument('--user_preference_gt', action='store_true',
                      help='Flag to use user preference ground-truth.')
  parser.add_argument('--hist_bins', type=int, default=48,
                      help='Number of bins in the histogram.')
  parser.add_argument('--device', type=str, choices=DEVICES, default='gpu',
                      help='Device to use for training. Options: "cpu" or "gpu".')
  parser.add_argument('--optimizer', type=str, choices=OPTIMIZERS, default='adam',
                      help='Optimization algorithm. Options are: "adam", "adamw", "nadam", "adamax", "sgd".')
  parser.add_argument('--exp_name', type=str, default='awb',
                      help='Experiment name.')

  return parser.parse_args()


def train(train_dataset: IlluminantEstimationDataLoader,
          validation_dataset: IlluminantEstimationDataLoader,
          training_config: Dict[str, Any]):
  """Trains an illuminant estimation model."""

  def lr_lambda(epoch):
    """Lambda function to increase learning rate during warm-up."""
    if epoch < warm_up_epochs:
      return (max_lr / initial_lr) ** (epoch / warm_up_epochs)
    else:
      return 1.0

  def perform_validation(model: IllumEstimator, data: DataLoader,
                         cluster: Optional[int]=None) -> float:
    """Performs validation of a given model."""
    model.eval()
    val_loss = 0
    i = 0
    for i, batch in enumerate(data):
      capture_data = batch['capture_data'].to(device=device_type)
      hist = batch['hist'].to(device=device_type)
      gt_illum = batch['gt'].to(device=device_type)
      if cluster is None:
        estimated_illum = model(hist, capture_data)
      else:
        estimated_illum = model(hist, capture_data, cluster)
      loss = loss_func(gt_illum, estimated_illum, shrink=False)
      val_loss += loss.item()

    if i == 0:
      return val_loss
    else:
      return val_loss / i

  epochs = training_config['epochs']
  lr = training_config['lr']
  weight_decay = training_config['weight_decay']
  device_type = training_config['device']
  output_path = training_config['output_path']
  exp_name = training_config['exp_name']
  batch_size = training_config['batch_size']
  optimizer = training_config['optimizer'].lower()


  if optimizer == 'adam':
    opt = optim.Adam
  elif optimizer == 'adamax':
    opt = optim.Adamax
  elif optimizer == 'nadam':
    opt = optim.NAdam
  elif optimizer == 'adamw':
    opt = optim.AdamW
  elif optimizer == 'sgd':
    opt = optim.SGD
  else:
    raise ValueError('Invalid optimizer name.')

  capture_data_size = 0
  train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True,
                            num_workers=NUM_WORKERS)
  val_loader = DataLoader(validation_dataset, batch_size=1, shuffle=False,
                          num_workers=NUM_WORKERS)
  for i, data in enumerate(train_loader):
    capture_data_size = data['capture_data'][0, ...].numel()
    if i == 0:
      break

  model = IllumEstimator(in_channels=capture_data_size,
                         hist_channels=4).to(device=device_type)

  model.print_num_of_params()
  loss_func = angular_loss


  initial_lr = 1e-6
  optimizer = opt(model.parameters(), lr=initial_lr, weight_decay=weight_decay)
  warm_up_epochs = 5
  max_lr = lr
  scheduler_annealing  = CosineAnnealingLR(optimizer, T_max=epochs - warm_up_epochs)
  scheduler_lambda = LambdaLR(optimizer, lr_lambda)


  training_config['device'] = None
  utils.write_json_file(training_config, os.path.join(output_path, f'config-{exp_name}.json'))

  best_valid_result = None

  for epoch in range(epochs):
    print(f'---------- Epoch {epoch + 1} -----------')
    if epoch % 100 == 0 and epoch != 0 and batch_size < MAXIMUM_BATCH_SIZE:
      print(f'Increasing batch size: {batch_size} -> {batch_size * 2}')
      batch_size = batch_size * 2
      train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True,
                                num_workers=NUM_WORKERS)
    for i, batch in enumerate(train_loader):
      capture_data = batch['capture_data'].to(device=device_type)
      hist = batch['hist'].to(device=device_type)
      gt_illum = batch['gt'].to(device=device_type)
      estimated_illum = model(hist, capture_data)
      loss = loss_func(gt_illum, estimated_illum, shrink=True)
      ae_loss = loss.item()
      optimizer.zero_grad()
      loss.backward()
      optimizer.step()
      if (i + 1) % 10 == 0:
        print(f'Epoch [{epoch + 1}/{epochs}], Step [{i + 1}/{len(train_loader)}], '
              f'AE Loss: {ae_loss:.4f}')

    if epoch < warm_up_epochs:
      scheduler_lambda.step()
    else:
      if epoch == warm_up_epochs:
        for param_group in optimizer.param_groups:
          param_group['lr'] = max_lr
      scheduler_annealing.step()

    val_loss = perform_validation(model, val_loader)
    model.train()
    print(
      f'Epoch [{epoch + 1}/{epochs}], Step [{i + 1}/{len(train_loader)}], AE Loss: {ae_loss:.4f}, '
      f'Val Loss: {val_loss:.4f}')
    if best_valid_result is None or val_loss < best_valid_result:
      torch.save(model.state_dict(), os.path.join(output_path, f'model-{exp_name}.pt'))
      best_valid_result = val_loss

  print(f'Finished training. Best validation loss: {best_valid_result}.')

if __name__ == '__main__':
  args = get_args()
  print_args_table(args)
  capture_data = args.capture_data
  capture_data = sorted(capture_data)
  if args.device == 'gpu':
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
  else:
    device = torch.device('cpu')

  if args.target_size is None:
    target_size = TARGET_SIZE
    print(f'Target image size is not given. It uses default size: {target_size}.')
  else:
    target_size = args.target_size

  training_config = {'epochs': args.epochs, 'lr': args.learning_rate, 'weight_decay': args.weight_decay,
                     'device': device, 'capture_data': capture_data, 'output_path': args.output_path,
                     'exp_name': args.exp_name, 'user_preference': args.user_preference_gt,
                     'hist_bins': args.hist_bins, 'batch_size': args.batch_size,
                     'optimizer': args.optimizer, 'target_size': target_size}

  print('Loading training data ...')
  tr_dataset = IlluminantEstimationDataLoader(
    data_dir=args.dataset_path, capture_data=capture_data, target_size=args.target_size,
    user_pref=args.user_preference_gt, train=True, hist_bins=args.hist_bins)

  norm_values = tr_dataset.get_normalization_values()
  training_config.update({'norm_values': norm_values})
  hist_boundaries = tr_dataset.get_hist_boundaries()
  training_config.update({'hist_boundaries': hist_boundaries})


  print('Loading validation data ...')
  val_dataset = IlluminantEstimationDataLoader(
    data_dir=args.dataset_path, capture_data=capture_data, target_size=args.target_size,
    valid=True, normalization_data=norm_values, user_pref=args.user_preference_gt,
    hist_bins=args.hist_bins, hist_boundaries=hist_boundaries)

  os.makedirs(args.output_path, exist_ok=True)

  print('Starting training ...')
  train(train_dataset=tr_dataset, validation_dataset=val_dataset, training_config=training_config)