image_to_image.model_interactions.train

  1"""
  2Module to train, validate, and evaluate image-to-image models. 
  3Supports multiple models, losses, optimizers, schedulers, mixed-precision training,
  4checkpointing, and experiment tracking with MLflow and TensorBoard.
  5
  6The train function handles full experiment orchestration including:
  7- Argument parsing and device setup.
  8- Dataset and dataloader initialization.
  9- Model, optimizer, loss function, and scheduler setup.
 10- Mixed precision (AMP) and warm-up handling.
 11- MLflow and TensorBoard logging.
 12- Periodic validation and checkpoint saving.
 13
 14Functions:
 15- get_loss
 16- get_optimizer
 17- get_scheduler
 18- backward_model
 19- train_one_epoch
 20- evaluate
 21- train
 22
 23By Tobia Ippolito
 24"""
 25# ---------------------------
 26#        > Imports <
 27# ---------------------------
 28import os
 29import shutil
 30import time
 31import copy
 32
 33import matplotlib.pyplot as plt
 34
 35from tqdm import tqdm
 36
 37import torch
 38from torch import nn, optim
 39from torch.utils.data import DataLoader
 40from torch.amp import autocast, GradScaler
 41import torchvision
 42
 43# Experiment tracking
 44import mlflow
 45import mlflow.pytorch
 46from torch.utils.tensorboard import SummaryWriter
 47
 48import prime_printer as prime
 49
 50
 51from ..utils.argument_parsing import parse_args
 52from ..utils.model_io import get_model, save_checkpoint
 53
 54from ..data.physgen import PhysGenDataset
 55from ..data.residual_physgen import PhysGenResidualDataset, to_device
 56
 57from ..models.resfcn import ResFCN
 58from ..models.pix2pix import Pix2Pix
 59from ..models.residual_design_model import ResidualDesignModel
 60from ..models.transformer import PhysicFormer
 61
 62from ..losses.weighted_combined_loss import WeightedCombinedLoss
 63from ..scheduler.warm_up import WarmUpScheduler
 64from..amp.dummy_scaler import DummyScaler
 65
 66
 67
 68# ---------------------------
 69#      > Train Helpers <
 70# ---------------------------
 71def get_data(args):
 72    if args.model.lower() == "residual_design_model":
 73        train_dataset = PhysGenResidualDataset(variation=args.data_variation, mode="train", 
 74                                               fake_rgb_output=args.fake_rgb_output, make_14_dividable_size=args.make_14_dividable_size,
 75                                               reflexion_channels=args.reflexion_channels, reflexion_steps=args.reflexion_steps, reflexions_as_channels=args.reflexions_as_channels)
 76        
 77        val_dataset = PhysGenDataset(variation=args.data_variation, mode="validation", input_type="osm", output_type="standard", 
 78                                    fake_rgb_output=args.fake_rgb_output, make_14_dividable_size=args.make_14_dividable_size,
 79                                    reflexion_channels=args.reflexion_channels, reflexion_steps=args.reflexion_steps, reflexions_as_channels=args.reflexions_as_channels)
 80    else:
 81        train_dataset = PhysGenDataset(variation=args.data_variation, mode="train", input_type=args.input_type, output_type=args.output_type, 
 82                                    fake_rgb_output=args.fake_rgb_output, make_14_dividable_size=args.make_14_dividable_size,
 83                                    reflexion_channels=args.reflexion_channels, reflexion_steps=args.reflexion_steps, reflexions_as_channels=args.reflexions_as_channels)
 84        
 85        val_dataset = PhysGenDataset(variation=args.data_variation, mode="validation", input_type=args.input_type, output_type=args.output_type, 
 86                                    fake_rgb_output=args.fake_rgb_output, make_14_dividable_size=args.make_14_dividable_size,
 87                                    reflexion_channels=args.reflexion_channels, reflexion_steps=args.reflexion_steps, reflexions_as_channels=args.reflexions_as_channels)
 88
 89    train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True)
 90    val_loader = DataLoader(val_dataset, batch_size=args.batch_size)
 91
 92    return train_dataset, val_dataset, train_loader, val_loader
 93
 94
 95
 96def get_loss(loss_name, args):
 97    """
 98    Returns a loss function instance based on the provided loss name.
 99
100    Supported losses:
101    - L1 / L1_2
102    - CrossEntropy / CrossEntropy_2
103    - WeightedCombined / WeightedCombined_2
104
105    Parameter:
106    - loss_name (str):
107        Name of the loss function.
108    - args:
109        Parsed command-line arguments with configured loss weights.
110
111    Returns:
112    - criterion (nn.Module): Instantiated loss function.
113    """
114    loss_name = loss_name.lower()
115
116    if loss_name == "l1":
117        criterion = nn.L1Loss()
118    elif loss_name == "l1_2":
119        criterion = nn.L1Loss()
120    elif loss_name == "crossentropy":
121        criterion = nn.CrossEntropyLoss()
122    elif loss_name == "crossentropy_2":
123        criterion = nn.CrossEntropyLoss()
124    elif loss_name == "weighted_combined":
125        criterion = WeightedCombinedLoss( 
126                        silog_lambda=args.wc_loss_silog_lambda, 
127                        weight_silog=args.wc_loss_weight_silog, 
128                        weight_grad=args.wc_loss_weight_grad,
129                        weight_ssim=args.wc_loss_weight_ssim,
130                        weight_edge_aware=args.wc_loss_weight_edge_aware,
131                        weight_l1=args.wc_loss_weight_l1,
132                        weight_var=args.wc_loss_weight_var,
133                        weight_range=args.wc_loss_weight_range,
134                        weight_blur=args.wc_loss_weight_blur
135                    )    
136    elif loss_name == "weighted_combined_2":
137        criterion = WeightedCombinedLoss( 
138                        silog_lambda=args.wc_loss_silog_lambda_2, 
139                        weight_silog=args.wc_loss_weight_silog_2, 
140                        weight_grad=args.wc_loss_weight_grad_2,
141                        weight_ssim=args.wc_loss_weight_ssim_2,
142                        weight_edge_aware=args.wc_loss_weight_edge_aware_2,
143                        weight_l1=args.wc_loss_weight_l1_2,
144                        weight_var=args.wc_loss_weight_var_2,
145                        weight_range=args.wc_loss_weight_range_2,
146                        weight_blur=args.wc_loss_weight_blur_2
147                    )    
148    else:
149        raise ValueError(f"'{loss_name}' is not a supported loss.")
150    
151    return criterion
152
153
154
155def get_optimizer(optimizer_name, model, lr, args):
156    """
157    Returns an optimizer for the given model.
158
159    Supported optimizers:
160    - Adam
161    - AdamW
162
163    Parameter:
164    - optimizer_name (str):
165        Name of the optimizer.
166    - model (nn.Module):
167        Model whose parameters should be optimized.
168    - lr (float):
169        Learning rate.
170    - args:
171        Parsed command-line arguments with optimizer configuration.
172
173    Returns:
174    - optimizer (torch.optim.Optimizer): Instantiated optimizer.
175    """
176    optimizer_name = optimizer_name.lower()
177
178    weight_decay_rate = args.weight_decay_rate if args.weight_decay else 0
179
180    if args.optimizer.lower() == "adam":
181        optimizer = optim.Adam(model.parameters(), lr=lr,  betas=(0.5, 0.999), weight_decay=weight_decay_rate)
182    elif args.optimizer.lower() == "adamw":
183        optimizer = optim.AdamW(model.parameters(), lr=lr,  betas=(0.5, 0.999), weight_decay=weight_decay_rate)
184    else:
185        raise ValueError(f"'{optimizer_name}' is not a supported optimizer.")
186    
187    return optimizer
188
189
190
191def get_scheduler(scheduler_name, optimizer, args):
192    """
193    Returns a learning rate scheduler for the given optimizer.
194
195    Supported schedulers:
196    - StepLR
197    - CosineAnnealingLR
198
199    Parameter:
200    - scheduler_name (str):
201        Name of the scheduler.
202    - optimizer:
203        Optimizer whose learning rate will be managed.
204    - args:
205        Parsed command-line arguments containing scheduler configuration.
206
207    Returns:
208    - scheduler (torch.optim.lr_scheduler): Instantiated scheduler.
209    """
210    scheduler_name = scheduler_name.lower()
211
212    if scheduler_name == "step":
213        scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.9)
214    elif scheduler_name == "cosine":
215        scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=args.epochs)
216    else:
217        raise ValueError(f"'{scheduler_name}' is not a supported scheduler.")
218    
219    return scheduler
220
221
222
223def get_params(args, device, n_model_params, current_save_name, checkpoint_save_dir):
224    return {
225            # General
226            "mode": args.mode,
227            "device": str(device),
228
229            # Training
230            "epochs": args.epochs,
231            "batch_size": args.batch_size,
232            "learning_rate": args.lr,
233            "loss_function": args.loss,
234            "optimizer": args.optimizer,
235            "weight_decay": args.weight_decay,
236            "weight_decay_rate": args.weight_decay_rate,
237            "gradient_clipping": args.gradient_clipping,
238            "gradient_clipping_threshold": args.gradient_clipping_threshold,
239            "scheduler": args.scheduler,
240            "use_warm_up": args.use_warm_up,
241            "warm_up_start_lr": args.warm_up_start_lr,
242            "warm_up_step_duration": args.warm_up_step_duration,
243            "use_amp": args.activate_amp,
244            "amp_scaler": args.amp_scaler,
245            "save_only_best_model": args.save_only_best_model,
246
247            # Loss
248            "wc_loss_silog_lambda": args.wc_loss_silog_lambda,
249            "wc_loss_weight_silog": args.wc_loss_weight_silog,
250            "wc_loss_weight_grad": args.wc_loss_weight_grad,
251            "wc_loss_weight_ssim": args.wc_loss_weight_ssim,
252            "wc_loss_weight_edge_aware": args.wc_loss_weight_edge_aware,
253            "wc_loss_weight_l1": args.wc_loss_weight_l1,
254            "wc_loss_weight_var": args.wc_loss_weight_var,
255            "wc_loss_weight_range": args.wc_loss_weight_range,
256            "wc_loss_weight_blur": args.wc_loss_weight_blur,
257
258            # Model
259            "model": args.model,
260            "n_model_params": n_model_params,
261            "resfcn_in_channels": args.resfcn_in_channels,
262            "resfcn_hidden_channels": args.resfcn_hidden_channels,
263            "resfcn_out_channels": args.resfcn_out_channels,
264            "resfcn_num_blocks": args.resfcn_num_blocks,
265
266            "pix2pix_in_channels": args.pix2pix_in_channels,
267            "pix2pix_hidden_channels": args.pix2pix_hidden_channels,
268            "pix2pix_out_channels": args.pix2pix_out_channels,
269            "pix2pix_second_loss_lambda": args.pix2pix_second_loss_lambda,
270
271            "physicsformer_in_channels": args.physicsformer_in_channels,
272            "physicsformer_out_channels": args.physicsformer_out_channels,
273            "physicsformer_img_size": args.physicsformer_img_size,
274            "physicsformer_patch_size": args.physicsformer_patch_size,
275            "physicsformer_embedded_dim": args.physicsformer_embedded_dim,
276            "physicsformer_num_blocks": args.physicsformer_num_blocks,
277            "physicsformer_heads": args.physicsformer_heads,
278            "physicsformer_mlp_dim": args.physicsformer_mlp_dim,
279            "physicsformer_dropout": args.physicsformer_dropout,
280
281            # Data
282            "data_variation": args.data_variation,
283            "input_type": args.input_type,
284            "output_type": args.output_type,
285            "fake_rgb_output": args.fake_rgb_output,
286            "make_14_dividable_size": args.make_14_dividable_size,
287
288            # Experiment tracking
289            "experiment_name": args.experiment_name,
290            "run_name": current_save_name, # CURRENT_SAVE_NAME,
291            "tensorboard_path": args.tensorboard_path,
292            "save_path": args.save_path,
293            "checkpoint_save_dir": checkpoint_save_dir,
294            "cmap": args.cmap,
295
296            # >> Residual Model <<
297            "base_model": args.base_model,
298            "complex_model": args.complex_model,
299            "combine_mode": args.combine_mode,
300
301            # ---- Loss (2nd branch)
302            "loss_2": args.loss_2,
303            "wc_loss_silog_lambda_2": args.wc_loss_silog_lambda_2,
304            "wc_loss_weight_silog_2": args.wc_loss_weight_silog_2,
305            "wc_loss_weight_grad_2": args.wc_loss_weight_grad_2,
306            "wc_loss_weight_ssim_2": args.wc_loss_weight_ssim_2,
307            "wc_loss_weight_edge_aware_2": args.wc_loss_weight_edge_aware_2,
308            "wc_loss_weight_l1_2": args.wc_loss_weight_l1_2,
309            "wc_loss_weight_var_2": args.wc_loss_weight_var_2,
310            "wc_loss_weight_range_2": args.wc_loss_weight_range_2,
311            "wc_loss_weight_blur_2": args.wc_loss_weight_blur_2,
312
313            # ---- ResFCN Model 2
314            "resfcn_2_in_channels": args.resfcn_2_in_channels,
315            "resfcn_2_hidden_channels": args.resfcn_2_hidden_channels,
316            "resfcn_2_out_channels": args.resfcn_2_out_channels,
317            "resfcn_2_num_blocks": args.resfcn_2_num_blocks,
318
319            # ---- Pix2Pix Model 2
320            "pix2pix_2_in_channels": args.pix2pix_2_in_channels,
321            "pix2pix_2_hidden_channels": args.pix2pix_2_hidden_channels,
322            "pix2pix_2_out_channels": args.pix2pix_2_out_channels,
323            "pix2pix_2_second_loss_lambda": args.pix2pix_2_second_loss_lambda,
324
325            # ---- PhysicsFormer Model 2
326            "physicsformer_in_channels_2": args.physicsformer_in_channels_2,
327            "physicsformer_out_channels_2": args.physicsformer_out_channels_2,
328            "physicsformer_img_size_2": args.physicsformer_img_size_2,
329            "physicsformer_patch_size_2": args.physicsformer_patch_size_2,
330            "physicsformer_embedded_dim_2": args.physicsformer_embedded_dim_2,
331            "physicsformer_num_blocks_2": args.physicsformer_num_blocks_2,
332            "physicsformer_heads_2": args.physicsformer_heads_2,
333            "physicsformer_mlp_dim_2": args.physicsformer_mlp_dim_2,
334            "physicsformer_dropout_2": args.physicsformer_dropout_2,
335        }
336
337
338def backward_model(model, x, y, optimizer, criterion, device, epoch, amp_scaler, gradient_clipping_threshold=None):
339    """
340    Performs a backward pass, optimizer step, and mixed-precision handling.
341
342    Parameter:
343    - model (nn.Module):
344        Model to train.
345    - x (torch.Tensor):
346        Input tensor.
347    - y (torch.Tensor):
348        Target tensor.
349    - optimizer:
350        Optimizer or tuple of optimizers (for Pix2Pix/ResidualDesignModel).
351    - criterion:
352        Loss function.
353    - device (torch.device):
354        Device to perform computation on.
355    - epoch (int):
356        Current training epoch.
357    - amp_scaler (GradScaler):
358        Gradient scaler for mixed-precision training.
359    - gradient_clipping_threshold (float, optional):
360        Maximum allowed gradient norm for clipping.
361
362    Returns:
363    - loss (torch.Tensor): Computed loss value for the batch.
364    """
365    if isinstance(model, Pix2Pix):
366        if epoch is None or epoch % 2 == 0:
367            model.discriminator_step(x, y, optimizer[1], amp_scaler, device, gradient_clipping_threshold)
368        loss, _, _ = model.generator_step(x, y, optimizer[0], amp_scaler, device, gradient_clipping_threshold)
369    elif amp_scaler:
370        # reset gradients 
371        optimizer.zero_grad()
372
373        with autocast(device_type=device.type):
374            y_predict = model(x)
375            loss = criterion(y_predict, y)
376        amp_scaler.scale(loss).backward()
377        if gradient_clipping_threshold:
378            # Unscale first!
379            amp_scaler.unscale_(optimizer)
380            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=gradient_clipping_threshold)
381        amp_scaler.step(optimizer)
382        amp_scaler.update()
383    else:
384        # reset gradients 
385        optimizer.zero_grad()
386
387        y_predict = model(x)
388        loss = criterion(y_predict, y)
389        loss.backward()
390        if gradient_clipping_threshold:
391            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=gradient_clipping_threshold)
392        optimizer.step()
393    return loss
394
395
396
397def train_one_epoch(model, loader, optimizer, criterion, device, epoch=None, amp_scaler=None, gradient_clipping_threshold=None):
398    """
399    Runs one full epoch of training and returns the average loss.
400
401    Parameter:
402    - model (nn.Module):
403        Model to train.
404    - loader (DataLoader):
405        Data loader containing training batches.
406    - optimizer:
407        Optimizer or tuple of optimizers.
408    - criterion:
409        Loss function or tuple of losses (for multi-stage models).
410    - device (torch.device):
411        Device to use for training.
412    - epoch (int, optional):
413        Current epoch index.
414    - amp_scaler (GradScaler, optional):
415        Mixed-precision scaler.
416    - gradient_clipping_threshold (float, optional):
417        Max allowed gradient norm.
418
419    Returns:
420    - avg_loss (float): Average training loss for the epoch.
421    """
422    # change to train mode -> calc gradients
423    model.train()
424
425    total_loss = 0.0
426    # for x, y in tqdm(loader, desc=f"Epoch {epoch:03}", leave=True, ascii=True mininterval=5):
427    for x, y in loader:
428
429        if not isinstance(model, ResidualDesignModel):
430            x, y = x.to(device), y.to(device)
431
432        if isinstance(model, ResidualDesignModel):
433            base_input, complex_input = x
434            base_target, complex_target, target_= y
435
436            # Basline
437            base_input = base_input.to(device)
438            base_target = base_target.to(device)
439            base_loss = backward_model(model=model.base_model, x=base_input, y=base_target, optimizer=optimizer[0], criterion=criterion[0], device=device, epoch=epoch, amp_scaler=amp_scaler, gradient_clipping_threshold=gradient_clipping_threshold)
440            # del base_input
441            # torch.cuda.empty_cache()
442
443            # Complex
444            complex_input = complex_input.to(device)
445            complex_target = complex_target.to(device)
446            complex_loss = backward_model(model=model.complex_model, x=complex_input, y=complex_target, optimizer=optimizer[1], criterion=criterion[1], device=device, epoch=epoch, amp_scaler=amp_scaler, gradient_clipping_threshold=gradient_clipping_threshold)
447            # del complex_input
448            # torch.cuda.empty_cache()
449
450            # Fusion
451            target_ = target_.to(device)
452            combine_loss = model.combine_net.backward(base_target, complex_target, target_)
453
454            model.backward(base_target, complex_target, target_)
455
456            model.last_base_loss = base_loss
457            model.last_complex_loss = complex_loss
458            model.last_combined_loss = combine_loss
459            loss = base_loss + complex_loss + combine_loss
460        else:
461            loss = backward_model(model=model, x=x, y=y, optimizer=optimizer, criterion=criterion, device=device, epoch=epoch, amp_scaler=amp_scaler, gradient_clipping_threshold=gradient_clipping_threshold)
462        total_loss += loss.item()
463    return total_loss / len(loader)
464
465
466
467@torch.no_grad()
468def evaluate(model, loader, criterion, device, writer=None, epoch=None, save_path=None, cmap="gray", use_tqdm=True):
469    """
470    Evaluates the model on the validation set and logs results to TensorBoard or MLflow.
471
472    Parameter:
473    - model (nn.Module):
474        Model to evaluate.
475    - loader (DataLoader):
476        Validation data loader.
477    - criterion:
478        Loss function used for evaluation.
479    - device (torch.device):
480        Device to perform inference on.
481    - writer (SummaryWriter, optional):
482        TensorBoard writer for visualization.
483    - epoch (int, optional):
484        Current epoch index for logging.
485    - save_path (str, optional):
486        Directory to save sample images.
487    - cmap (str):
488        Colormap for saved images.
489    - use_tqdm (bool):
490        Whether to use tqdm progress bar.
491
492    Returns:
493    - avg_loss (float): Average validation loss.
494    """
495    model.eval()
496
497    total_loss = 0.0
498    is_first_round = True
499
500    if use_tqdm:
501        validation_iter = tqdm(loader, desc="Validation", ascii=True, mininterval=3)
502    else:
503        validation_iter = loader
504
505    for x, y in validation_iter:
506        x, y = x.to(device), y.to(device)
507        y_predict = model(x)
508        total_loss += criterion(y_predict, y).item()
509
510        if is_first_round:
511            if writer:
512                # Convert to grid
513                img_grid_input = torchvision.utils.make_grid(x[:4].cpu(), normalize=True, scale_each=True)
514                max_val = max(y_predict.max().item(), 1e-8)
515                if y_predict.ndim == 3:  # e.g., [B, H, W]
516                    img_grid_pred = torchvision.utils.make_grid(y_predict[:4].unsqueeze(1).float().cpu() / max_val)
517                else:  # [B, 1, H, W]
518                    img_grid_pred = torchvision.utils.make_grid(y_predict[:4].float().cpu() / max_val)
519
520                max_val = max(y.max().item(), 1e-8)
521                if y.ndim == 3:
522                    img_grid_gt = torchvision.utils.make_grid(y[:4].unsqueeze(1).float().cpu() / max_val)
523                else:
524                    img_grid_gt = torchvision.utils.make_grid(y[:4].float().cpu() / max_val)
525
526                # Log to TensorBoard
527                writer.add_image("Input", img_grid_input, epoch)
528                writer.add_image("Prediction", img_grid_pred, epoch)
529                writer.add_image("GroundTruth", img_grid_gt, epoch)
530
531            if save_path:
532                # os.makedirs(save_path, exist_ok=True)
533                prediction_path = os.path.join(save_path, f"{epoch}_prediction.png")
534                plt.imsave(prediction_path, 
535                        y_predict[0].detach().cpu().numpy().squeeze(), cmap=cmap)
536                mlflow.log_artifact(prediction_path, artifact_path="images")
537
538                input_path = os.path.join(save_path, f"{epoch}_input.png")
539                plt.imsave(input_path, 
540                        x[0][0].detach().cpu().numpy().squeeze(), cmap=cmap)
541                mlflow.log_artifact(input_path, artifact_path="images")
542
543                ground_truth_path = os.path.join(save_path, f"{epoch}_ground_truth.png")
544                plt.imsave(ground_truth_path, 
545                        y[0].detach().cpu().numpy().squeeze(), cmap=cmap)
546                mlflow.log_artifacts(ground_truth_path, artifact_path="images")
547
548                # alternative direct save:
549                # import numpy as np
550                # import io
551                # from PIL import Image
552
553                # img = (y[0].detach().cpu().squeeze().numpy() * 255).astype(np.uint8)
554                # img_pil = Image.fromarray(img)
555
556                # buf = io.BytesIO()
557                # img_pil.save(buf, format='PNG')
558                # buf.seek(0)
559
560                # mlflow.log_image(image=img_pil, artifact_file=f"images/{epoch}_ground_truth.png")
561                
562        is_first_round = False
563
564    return total_loss / len(loader)
565
566
567
568# ---------------------------
569#        > Train Main <
570# ---------------------------
571def train(args=None):
572    """
573    Main training loop for image-to-image tasks.
574
575    Workflow:
576    1. Initializes the training and validation datasets based on model type.
577    2. Constructs the model and its loss functions.
578    3. Configures optimizers, learning rate schedulers, and optional warm-up phases.
579    4. Enables mixed precision (AMP) if selected.
580    5. Sets up MLflow experiment tracking and TensorBoard visualization.
581    6. Executes the epoch loop:
582        - Trains the model for one epoch (`train_one_epoch()`).
583        - Optionally evaluates on the validation set.
584        - Logs metrics and learning rates.
585        - Updates the scheduler.
586        - Saves checkpoints (best or periodic).
587    7. Logs the trained model and experiment results to MLflow upon completion.
588
589    Parameters:
590    - args : argparse.Namespace, optional
591        Parsed command-line arguments containing all training configurations.
592        If None, the function will automatically call `parse_args()` to obtain them.
593    
594    Returns:
595    - None: The function performs training and logging in-place without returning values.
596
597    Notes:
598    - Automatically handles model-specific configurations (e.g., Pix2Pix discriminator, ResidualDesignModel branches).
599    - Uses `prime.get_time()` to generate time-stamped run names.
600    - Supports gradient clipping and various learning rate schedulers.
601
602    Logging:
603    - **MLflow**: Stores metrics, hyperparameters, checkpoints, and final model.
604    - **TensorBoard**: Logs training/validation losses, learning rates, and sub-loss components.
605    """
606    print("\n---> Welcome to Image-to-Image Training <---")
607
608    print("\nChecking your Hardware:")
609    print(prime.get_hardware())
610
611    # Parse arguments
612    if args is None:
613        args = parse_args()
614
615    CURRENT_SAVE_NAME = prime.get_time(pattern="YEAR-MONTH-DAY_HOUR_MINUTE_", time_zone="Europe/Berlin") + args.run_name
616
617    device = torch.device(args.device if torch.cuda.is_available() else "cpu")
618
619    # Dataset loading
620    train_dataset, val_dataset, train_loader, val_loader = get_data(args)
621
622    # Loss
623    criterion = get_loss(loss_name=args.loss, args=args)
624    
625    if args.model.lower() == "residual_design_model":
626        criterion = [
627            criterion,
628            get_loss(loss_name=args.loss_2+"_2", args=args)
629        ]
630
631    # Model Loading
632    model = get_model(args, device, criterion=criterion)
633
634    # get parameter amount
635    n_model_params = 0
636    for cur_model_param in model.parameters():
637        n_model_params += cur_model_param.numel()
638
639    INPUT_CHANNELS = model.get_input_channels()
640
641    # Optimizer
642    if args.model.lower() == "pix2pix":
643        optimizer = [get_optimizer(optimizer_name=args.optimizer, model=model.generator, lr=args.lr, args=args),
644                     get_optimizer(optimizer_name=args.optimizer_2, model=model.discriminator, lr=args.lr, args=args)]
645    elif args.model.lower() == "residual_design_model":
646        optimizer = [get_optimizer(optimizer_name=args.optimizer, model=model.base_model, lr=args.lr, args=args),
647                     get_optimizer(optimizer_name=args.optimizer_2, model=model.complex_model, lr=args.lr, args=args)]
648        
649        if args.base_model.lower() == "pix2pix":
650            optimizer[0] = [get_optimizer(optimizer_name=args.optimizer, model=model.base_model.generator, lr=args.lr, args=args), 
651                            get_optimizer(optimizer_name=args.optimizer, model=model.base_model.discriminator, lr=args.lr, args=args)]
652        if args.complex_model.lower() == "pix2pix":
653            optimizer[1] = [get_optimizer(optimizer_name=args.optimizer_2, model=model.complex_model.generator, lr=args.lr, args=args),
654                            get_optimizer(optimizer_name=args.optimizer_2, model=model.complex_model.discriminator, lr=args.lr, args=args)]
655    else:
656        optimizer = get_optimizer(optimizer_name=args.optimizer, model=model, lr=args.lr, args=args)
657
658    # Scheduler
659    if args.model.lower() == "residual_design_model":
660
661        if args.base_model.lower() == "pix2pix":
662            scheduler_1 = [get_scheduler(scheduler_name=args.scheduler, optimizer=optimizer[0][0], args=args),
663                           get_scheduler(scheduler_name=args.scheduler, optimizer=optimizer[0][1], args=args)]
664        else:
665            scheduler_1 = get_scheduler(scheduler_name=args.scheduler, optimizer=optimizer[0], args=args)
666            
667        if args.complex_model.lower() == "pix2pix":
668            scheduler_2 = [get_scheduler(scheduler_name=args.scheduler_2, optimizer=optimizer[1][0], args=args),
669                           get_scheduler(scheduler_name=args.scheduler_2, optimizer=optimizer[1][1], args=args)]
670        else:
671            scheduler_2 = get_scheduler(scheduler_name=args.scheduler_2, optimizer=optimizer[1], args=args)
672            
673            scheduler = [scheduler_1, scheduler_2]
674    elif args.model.lower() == "pix2pix":
675        scheduler = [get_scheduler(scheduler_name=args.scheduler, optimizer=optimizer[0], args=args),
676                     get_scheduler(scheduler_name=args.scheduler_2, optimizer=optimizer[1], args=args)]
677    else:
678        scheduler = get_scheduler(scheduler_name=args.scheduler, optimizer=optimizer, args=args)
679
680    # Warm-Up Scheduler
681    if args.use_warm_up:
682        if isinstance(scheduler, (tuple, list)):
683            new_scheduler = []
684            for cur_scheduler in scheduler:
685                new_scheduler += [WarmUpScheduler(start_lr=args.warm_up_start_lr, end_lr=args.lr, optimizer=cur_scheduler.optimizer, scheduler=cur_scheduler, step_duration=args.warm_up_step_duration)]
686
687            scheduler = new_scheduler
688        else:
689            if scheduler is not None:
690                cur_optimizer = scheduler.optimizer
691            else:
692                cur_optimizer = None
693            scheduler = WarmUpScheduler(start_lr=args.warm_up_start_lr, end_lr=args.lr, optimizer=cur_optimizer, scheduler=scheduler, step_duration=args.warm_up_step_duration)
694
695    # AMP Scaler
696    if args.activate_amp == False:
697        amp_scaler = None
698    elif args.amp_scaler.lower() == "none":
699        amp_scaler = DummyScaler()
700    elif args.amp_scaler.lower() == "grad":
701        amp_scaler = GradScaler()
702    else:
703        raise ValueError(f"'{args.amp_scaler}' is not an supported scaler.")
704
705    # setup checkpoint saving
706    checkpoint_save_dir = os.path.join(args.checkpoint_save_dir, args.experiment_name, CURRENT_SAVE_NAME)
707    os.makedirs(checkpoint_save_dir, exist_ok=True)
708    shutil.rmtree(checkpoint_save_dir)
709    os.makedirs(checkpoint_save_dir, exist_ok=True)
710
711    # setup intermediate image saving
712    save_path = os.path.join(args.save_path, args.experiment_name, CURRENT_SAVE_NAME)
713    os.makedirs(save_path, exist_ok=True)
714    shutil.rmtree(save_path)
715    os.makedirs(save_path, exist_ok=True)
716
717    # setup gradient clipping
718    if args.gradient_clipping:
719        gradient_clipping_threshold = args.gradient_clipping_threshold
720    else:
721        gradient_clipping_threshold = None
722
723    mlflow.set_experiment(args.experiment_name)
724    # same as:
725        # mlflow.create_experiment(args.experiment_name)
726        # mlflow.get_experiment_by_name(experiment_name)
727
728    # Start MLflow run
729    with mlflow.start_run(run_name=CURRENT_SAVE_NAME):
730
731        # TensorBoard writer
732        tensorboard_path = os.path.join(args.tensorboard_path, args.experiment_name, CURRENT_SAVE_NAME)
733        os.makedirs(tensorboard_path, exist_ok=True)
734        shutil.rmtree(tensorboard_path)
735        os.makedirs(tensorboard_path, exist_ok=True)
736        writer = SummaryWriter(log_dir=tensorboard_path)
737
738        # Log hyperparameters
739        params =  get_params(args=args, 
740                             device=device, 
741                             n_model_params=n_model_params, 
742                             current_save_name=CURRENT_SAVE_NAME, 
743                             checkpoint_save_dir=checkpoint_save_dir)
744        
745        # mlflow.log_params(vars(args))
746        mlflow.log_params(params)
747
748        params_text = "\n".join([f"{k}: {v}" for k, v in params.items()])
749        writer.add_text("hyperparameters", params_text, 0)
750
751        print(f"Train dataset size: {len(train_dataset)} | Validation dataset size: {len(val_dataset)}")
752        mlflow.log_metrics({
753            "train_dataset_size": len(train_dataset),
754            "val_dataset_size": len(val_dataset)
755        })
756
757
758        # log model architecture
759        #    -> create dataset for that
760        if isinstance(INPUT_CHANNELS, (list, tuple)):
761            dummy_inference_data = [torch.ones(size=(1, channels, 256, 256)).to(device) for channels in INPUT_CHANNELS]
762        else:
763            dummy_inference_data = torch.ones(size=(1, INPUT_CHANNELS, 256, 256)).to(device)
764        writer.add_graph(model, dummy_inference_data)
765
766        # Run Training
767        last_best_loss = float("inf")
768        try:
769            epoch_iter = tqdm(range(1, args.epochs + 1), desc="Epochs", ascii=True)
770            for epoch in epoch_iter:
771                start_time = time.time()
772                train_loss = train_one_epoch(model, train_loader, optimizer, criterion, device, epoch, amp_scaler, gradient_clipping_threshold)
773                duration = time.time() - start_time
774                if epoch % args.validation_interval == 0:
775                    val_loss = evaluate(model, val_loader, criterion, device, writer=writer, epoch=epoch, save_path=save_path, cmap=args.cmap, use_tqdm=False)
776                else:
777                    val_loss = float("inf")
778
779                val_str = f"{val_loss:.4f}" if epoch % args.validation_interval == 0 else "N/A"
780                epoch_iter.set_postfix(train_loss=f"{train_loss:.4f}", val_loss=val_str, time_needed=f"{duration:.2f}s")
781                # tqdm.write(f" -> Train Loss: {train_loss:.4f} | Val Loss: {val_str} | Time: {duration:.2f}")
782                # \n\n[Epoch {epoch:02}/{args.epochs}]
783                
784                # Hint: Tensorboard and mlflow does not like spaces in tags!
785
786                # Log to TensorBoard
787                writer.add_scalar("Time/epoch_duration", duration, epoch)
788                writer.add_scalar("Loss/train", train_loss, epoch)
789                writer.add_scalar("Loss/val", val_loss, epoch)
790                if isinstance(scheduler, list):
791                    if isinstance(scheduler[0], list):
792                        writer.add_scalar("LR/generator", scheduler[0][0].get_last_lr()[0], epoch)
793                        writer.add_scalar("LR/discriminator", scheduler[0][1].get_last_lr()[0], epoch)
794                    else:
795                        name = "generator" if args.model.lower() == "pix2pix" else "base_model"
796                        writer.add_scalar(f"LR/{name}", scheduler[0].get_last_lr()[0], epoch)
797
798                    if isinstance(scheduler[1], list):
799                        writer.add_scalar("LR/generator", scheduler[1][0].get_last_lr()[0], epoch)
800                        writer.add_scalar("LR/discriminator", scheduler[1][1].get_last_lr()[0], epoch)
801                    else:
802                        name = "discriminator" if args.model.lower() == "pix2pix" else "complex_model"
803                        writer.add_scalar(f"LR/{name}", scheduler[1].get_last_lr()[0], epoch)
804                else:
805                    writer.add_scalar("LR", scheduler.get_last_lr()[0], epoch)
806
807                # Log to MLflow
808                if type(scheduler) in [list, tuple]:
809                    metrics = {
810                        "train_loss": train_loss,
811                        "val_loss": val_loss,
812                    }
813                    idx = 0
814                    for idx, cur_scheduler in enumerate(scheduler):
815                        metrics[f"lr_{idx}"] = cur_scheduler.get_last_lr()[0]
816                    mlflow.log_metrics(metrics, step=epoch)
817                else:
818                    mlflow.log_metrics({
819                        "train_loss": train_loss,
820                        "val_loss": val_loss,
821                        "lr": scheduler.get_last_lr()[0]
822                    }, step=epoch)
823
824                # add sub losses / loss components
825                if args.model.lower() in ["pix2pix", "residual_design_model"]:
826                    losses = model.get_dict()
827                    for name, value in losses.items():
828                        writer.add_scalar(f"LossComponents/{name}", value, epoch)
829                    mlflow.log_metrics(losses, step=epoch)
830
831                if type(criterion) in [list, tuple]:
832                    if args.loss in ["weighted_combined"]:
833                        losses = criterion[0].get_dict()
834                        for name, value in losses.items():
835                            writer.add_scalar(f"LossComponents/{name}", value, epoch)
836                        mlflow.log_metrics(losses, step=epoch)
837
838                    if args.loss_2 in ["weighted_combined"] and args.model.lower() in ["residual_design_model"]:
839                        losses = criterion[1].get_dict()
840                        for name, value in losses.items():
841                            writer.add_scalar(f"LossComponents/{name}", value, epoch)
842                        mlflow.log_metrics(losses, step=epoch)
843                else:
844                    if args.loss in ["weighted_combined"]:
845                        losses = criterion.get_dict()
846                        for name, value in losses.items():
847                            writer.add_scalar(f"LossComponents/{name}", value, epoch)
848                        mlflow.log_metrics(losses, step=epoch)
849
850                
851
852                # Step scheduler
853                if args.model.lower() in ["pix2pix"]:
854                    scheduler[0].step()
855                    scheduler[1].step()
856                else:
857                    scheduler.step()
858
859                # Save Checkpoint
860                if args.save_only_best_model:
861                    if val_loss < last_best_loss or (last_best_loss == float("inf") and epoch == args.epochs):
862                        last_best_loss = val_loss
863                        checkpoint_path = os.path.join(checkpoint_save_dir, f"best_checkpoint.pth")
864                        save_checkpoint(args, model, optimizer, scheduler, epoch, checkpoint_path)
865
866                        # Log model checkpoint path
867                        mlflow.log_artifact(checkpoint_path)
868                elif epoch % args.checkpoint_interval == 0 or epoch == args.epochs:
869                    checkpoint_path = os.path.join(checkpoint_save_dir, f"epoch_{epoch}.pth")
870                    save_checkpoint(args, model, optimizer, scheduler, epoch, checkpoint_path)
871
872                    # Log model checkpoint path
873                    mlflow.log_artifact(checkpoint_path)
874
875            # Log final model
876            try:
877                mlflow.pytorch.log_model(model.cpu(), name="model", input_example=dummy_inference_data.cpu().numpy())
878            except Exception as e:
879                print(e)
880                mlflow.pytorch.log_model(model, name="model")
881            mlflow.end_run()
882        finally:
883            writer.close()
884
885        print("Training completed.")
886
887
888
889if __name__ == "__main__":
890    train()