implemented loss balancer

lib/infer_pack/commons.py

@@ -2,20 +2,71 @@ import math
 import torch
 from torch.nn import functional as F
 
-def median_pool1d(signal, kernel_size, dim=-1,stride=1):
+def median_pool1d(signal, kernel_size, dim=-1, stride=1):
+    """
+    Apply a 1D median filter along a specified dimension for a 3D tensor.
+
+    Args:
+        signal (torch.Tensor): The input tensor of shape (batch_size, F, T).
+        kernel_size (int): The size of the median filter kernel.
+        dim (int): The dimension along which to apply the median filter.
+        stride (int): The stride of the sliding window.
+
+    Returns:
+        torch.Tensor: The filtered tensor.
+    """
+    # Ensure the input tensor is 3D (batch_size, F, T)
+    if signal.dim() != 3:
+        raise ValueError("Input tensor must be 3D (batch_size, F, T)")
+
+    # Adjust kernel size if larger than input dimension
+    input_size = signal.size(dim)
+    if kernel_size > input_size:
+        kernel_size = input_size
+
     # Pad the input to handle edge cases
     padding = (kernel_size - 1) // 2
+    signal = F.pad(signal, (padding, padding), mode='reflect')
+    if dim == -1:
+        signal = F.pad(signal, (padding, padding), mode='reflect')
+    elif dim == -2:
+        signal = F.pad(signal, (0, 0, padding, padding), mode="reflect")
 
-    input_padded = F.pad(signal, (padding, padding), mode='reflect')
-    
     # Unfold the input to create sliding windows
-    unfolded = input_padded.unfold(-1, kernel_size, stride)
-    
+    unfolded = signal.unfold(dim, kernel_size, stride)
+
     # Compute the median along the last dimension
-    median = unfolded.median(dim=dim).values
-    
+    median = unfolded.median(dim=-1).values
+
     return median
 
+def autocorrelation1d(log_magnitude, max_lag=None):
+    """
+    Computes the autocorrelation of log-magnitude spectrogram.
+    :param log_magnitude: Log magnitude spectrogram (batch, channels, time)
+    :param max_lag: Maximum lag for which to compute autocorrelation.
+    :return: Autocorrelation of the log-magnitude spectrogram.
+    """
+    if max_lag is None:
+        max_lag = log_magnitude.size(-1) // 2  # Default to half the signal length
+
+    # Ensure the signal is zero-mean
+    log_magnitude = log_magnitude - log_magnitude.mean(dim=-1, keepdim=True)
+    
+    # Compute autocorrelation using FFT-based convolution
+    padded_log_magnitude = F.pad(log_magnitude, (max_lag, 0), mode='reflect')
+    
+    # Ensure the kernel has the correct shape for conv1d
+    kernel = log_magnitude
+    padded_log_magnitude = padded_log_magnitude
+    
+    autocorr = F.conv1d(padded_log_magnitude, kernel, groups=log_magnitude.size(1))
+    
+    # Take only the relevant part
+    autocorr = F.normalize(autocorr[:,:max_lag + 1], p=2, dim=-1)
+    
+    return autocorr.nan_to_num(0)
+
 def init_weights(m, mean=0.0, std=0.01):
     classname = m.__class__.__name__
     if classname.find("Conv") != -1:
@@ -179,3 +230,7 @@ def clip_grad_value_(parameters, clip_value, norm_type=2, batch_size=1):
             p.grad.data.clamp_(min=-clip_value, max=clip_value)
     total_norm = total_norm ** (1.0 / norm_type)
     return total_norm
+
+def serialize_tensor(tensor):
+    if hasattr(tensor,"item"): return tensor.item()
+    else: return float(f"{float(tensor):.3f}")

lib/train/losses.py

@@ -1,30 +1,179 @@
 import torch
 import torchaudio
 import torch.nn.functional as F
-from ..infer_pack.commons import median_pool1d
 
-def log1p_autocorrelation(log_magnitude, max_lag=None):
-    """
-    Computes the autocorrelation of log-magnitude spectrogram.
-    :param log_magnitude: Log magnitude spectrogram (batch, channels, time)
-    :param max_lag: Maximum lag for which to compute autocorrelation.
-    :return: Autocorrelation of the log-magnitude spectrogram.
-    """
-    if max_lag is None:
-        max_lag = log_magnitude.size(-1) // 2  # Default to half the signal length
+from ..utils import gc_collect
+from ..infer_pack.commons import median_pool1d, autocorrelation1d
 
-    # Ensure the signal is zero-mean
-    log_magnitude = log_magnitude - log_magnitude.mean(dim=-1, keepdim=True)
+class LossBalancer:
+    model: torch.nn.Module
+
+    def __init__(self, model, initial_weights={}, historical_losses={}, ema_weights={}, epsilon=1e-8, weights_decay=0., loss_decay=.0, active=True, use_partial=False, use_pareto=True, use_norm=False):
+        """
+        Initializes the LossBalancer with optional initial weights, partial gradient computation, and EMA.
+        
+        Args:
+            model (torch.nn.Module): The model for which gradients will be computed.
+            initial_weights (dict): Initial weights for each loss term.
+            epsilon (float): A small value to avoid division by zero.
+            partial (bool): Whether to compute gradients only for the output layer.
+        """
+        self.model = model
+        self.epsilon = epsilon
+        self.use_partial = use_partial
+        self.weights_decay = weights_decay
+        self.loss_decay = loss_decay
+        self.initial_weights = initial_weights
+        self.ema_weights = ema_weights
+        self.historical_losses = historical_losses
+        self.active = active
+        self.use_pareto = use_pareto # apply the 20/80 rule to loss balancing
+        self.use_norm = use_norm
+
+
+    def to_dict(self):
+        return dict(
+            epsilon = self.epsilon,
+            use_partial = self.use_partial,
+            weights_decay = self.weights_decay,
+            loss_decay = self.loss_decay,
+            ema_weights = self.ema_weights,
+            initial_weights = self.initial_weights,
+            historical_losses = self.historical_losses,
+            active = self.active,
+            use_pareto = self.use_pareto, # apply the 20/80 rule to loss balancing
+            use_norm = self.use_norm
+        )
+
+    def update_ema_weights(self, new_weights):
+        """Updates the EMA of the weights."""
+        if not self.ema_weights: self.ema_weights = new_weights
+        else: self.ema_weights = {
+                k: self.weights_decay * self.ema_weights.get(k, 0.) + (1 - self.weights_decay) * new_weights[k]
+                for k in new_weights
+            }
     
-    # Compute autocorrelation using FFT-based convolution
-    padded_log_magnitude = F.pad(log_magnitude, (0, max_lag))
-    # autocorr = torch.corrcoef(padded_log_magnitude)
-    autocorr = F.conv1d(padded_log_magnitude, log_magnitude.flip(dims=[-1]))
+    def update_historical_losses(self, new_losses):
+        """Updates the EMA of the weights."""
+        if not self.historical_losses: self.historical_losses = new_losses
+        else: self.historical_losses = {
+                k: self.loss_decay * self.historical_losses.get(k, 0.) + (1 - self.loss_decay) * new_losses[k]
+                for k in new_losses
+            }
+
+    def calculate_loss_slope(self, key, current_loss):
+        """Calculates the slope of the loss using the current loss and its EMA."""
+        ema_loss = self.historical_losses.get(key, current_loss)+self.epsilon
+        slope = abs((current_loss-ema_loss)/ema_loss)  # relative loss change
+        if slope>1: slope*=slope # punish large changes
+        return slope
     
-    # Take only the relevant part and normalize
-    autocorr = F.normalize(autocorr[:, :max_lag + 1],p=2,dim=-1)
-    
-    return autocorr.nan_to_num(0)
+    def calculate_gradients(self, current_loss):
+        """Calculates the gradient norm of the current loss wrt the model params."""
+        self.model.zero_grad()  # Clear previous gradients
+
+        # Backward pass with partial or full gradient computation
+        model_params = list(self.model.parameters())[-1] if self.use_partial else self.model.parameters()
+        output_params = torch.autograd.grad([current_loss], [model_params], retain_graph=True, only_inputs=True)
+
+        # Compute L2 gradient norm
+        grad_norm = torch.stack([param.norm() for param in output_params if hasattr(param, "norm")]).norm(2)
+        del output_params
+        return grad_norm
+
+    def pareto_normalizer(self, data: dict):
+        # Sort the data in descending order based on the values
+        sorted_data = sorted(data.items(), key=lambda x: x[1], reverse=True)
+        
+        # Calculate the number of elements in the top 20%
+        top_20_count = max(int(0.2 * len(sorted_data)),1)
+        
+        # Calculate the weights
+        weights = [0] * len(sorted_data)
+        total_weight = 1.0
+        top_20_weight = 0.8 * total_weight
+        remaining_weight = total_weight - top_20_weight
+        
+        # Assign weights to the top 20%
+        for i in range(top_20_count):
+            weights[i] = top_20_weight / top_20_count
+        
+        # Assign weights to the remaining 80%
+        for i in range(top_20_count, len(sorted_data)):
+            weights[i] = remaining_weight / (len(sorted_data) - top_20_count)
+        
+        # Normalize the weights
+        total_weight = sum(weights)
+        normalized_weights = [w / total_weight for w in weights]
+        
+        # Combine the keys with their normalized weights
+        normalized_data = {sorted_data[i][0]: normalized_weights[i] for i in range(len(sorted_data))}
+        
+        return normalized_data
+
+    def on_train_batch_start(self, losses: dict):
+        """
+        Balances the loss terms based on their gradients before each training batch.
+        
+        Args:
+            losses (dict): Dictionary of individual loss terms (scalar tensors).
+        
+        Returns:
+            balanced_loss (torch.Tensor): Weighted sum of losses.
+            new_weights (dict): Updated weights for each loss term.
+        """
+        
+        if not self.initial_weights: self.initial_weights = {k: 1.0 for k in losses} # initialize weights
+        if not self.active: return sum(v*self.initial_weights[k] if k in self.initial_weights else v for k,v in losses)
+
+        gradients = {}
+        valid_losses = {}
+
+        # Process each loss, skip if the corresponding weight is 0
+        for key, loss in losses.items():
+            weight = self.initial_weights[key] if key in self.initial_weights else 1.
+            if weight == 0 or loss == 0 or not isinstance(loss,torch.Tensor): continue  # Skip loss with weight 0 or constants
+
+            if self.use_norm:
+                # Compute L2 gradient norm
+                grad_norm = self.calculate_gradients(loss)
+                gradients[key] = grad_norm.item()
+            else:
+                # Use loss plateau detection with EMA
+                loss_slope = self.calculate_loss_slope(key, loss)
+                gradients[key] = loss_slope.item()  # Smaller slope -> plateau -> lower priority
+            valid_losses[key] = loss
+
+        if not valid_losses: return torch.tensor(0.0)  # If all losses are skipped
+
+        # Calculate loss weights based on gradient magnitudes
+        if self.use_pareto: #20% hardest tasks gets 80% weight
+            normalized_weights = self.pareto_normalizer(gradients)
+        else:
+            total_gradient = sum(gradients.values()) + self.epsilon
+            normalized_weights = {k: w/total_gradient for k, w in gradients.items()}
+
+        # Update EMA weights
+        self.update_ema_weights(normalized_weights)
+        balanced_loss = sum(self.ema_weights[k] * loss for k, loss in valid_losses.items())
+        
+        # update historical losses
+        self.update_historical_losses({k: v.item() for k,v in valid_losses.items()})
+        
+        return balanced_loss
+
+    def on_epoch_end(self, weights_decay=None, loss_decay=None):
+        """
+        Optional method to call at the end of each epoch for logging or further adjustments.
+        """
+        if weights_decay is not None: self.weights_decay=weights_decay
+        if loss_decay is not None: self.loss_decay=loss_decay
+        print(f"|| EMA loss weights: {self.ema_weights} ({self.weights_decay=})")
+        print(f"|| EMA losses: {self.historical_losses} ({self.loss_decay=})")
+        weighted_loss = sum(v*self.ema_weights.get(k,0.) for k,v in self.historical_losses.items())
+        print(f"===> EMA loss: {weighted_loss:.3f} <====")
+        gc_collect()
+
 
 def compute_phase_loss(original_log_magnitude, generated_log_magnitude, max_lag=None):
     """
@@ -35,79 +184,78 @@ def compute_phase_loss(original_log_magnitude, generated_log_magnitude, max_lag=
     :return: Loss value.
     """
     # Compute autocorrelations
-    original_autocorr = log1p_autocorrelation(original_log_magnitude, max_lag)
-    generated_autocorr = log1p_autocorrelation(generated_log_magnitude, max_lag)
+    original_autocorr = autocorrelation1d(original_log_magnitude, max_lag)
+    generated_autocorr = autocorrelation1d(generated_log_magnitude, max_lag)
 
     # Compute the loss between autocorrelations
     loss = F.smooth_l1_loss(generated_autocorr, original_autocorr)
     
     return loss
-    
-def compute_temporal_envelope(log_magnitude: torch.Tensor, kernel_size=31):
-    """Compute the temporal envelope loss using median pooling and second-order differences."""
-    
-    # Apply median pooling to the log-magnitude spectrogram
-    pooled_stft = median_pool1d(log_magnitude, kernel_size, stride=kernel_size // 2 + 1)
-    
-    # Compute second-order differences on the median-pooled audio
-    diff = torch.diff(pooled_stft, dim=-1, n=2)
-    
-    return diff.nan_to_num(0)
 
-def compute_cepstrals(log_magnitude: torch.Tensor, kernel_size=31):
+def compute_envelope(log_magnitude: torch.Tensor, dim=-1, kernel_size=3):
 
-    pooled_magnitude = median_pool1d(log_magnitude, kernel_size, stride=kernel_size//2+1)
+    # average over axis
+    pooled_magnitude = log_magnitude.mean(dim)
 
-    # Compute the cepstrum
-    cepstral_coefficients = torch.fft.ifft(pooled_magnitude, dim=-2).abs()
+    # use max pool to get the peaks
+    max_filtered_tensor = F.max_pool1d(pooled_magnitude, kernel_size=kernel_size, stride=1, padding=kernel_size//2)
     
-    return cepstral_coefficients.sum(dim=-2).nan_to_num(0)
+    return max_filtered_tensor.nan_to_num(0).sum(-2)
 
-def compute_sts_loss(gen_stft: torch.Tensor, org_stft: torch.Tensor, kernel_size=31):
+def compute_sts_loss(gen_stft: torch.Tensor, org_stft: torch.Tensor,eps=1e-6):
     
     # B?N?T
-    org_mag = torch.log1p(org_stft.abs())
-    gen_mag = torch.log1p(gen_stft.abs())
-    org_temp = torch.log1p(org_stft.abs().sum(-2,keepdim=True))
-    gen_temp = torch.log1p(gen_stft.abs().sum(-2,keepdim=True))
+    org_mag = torch.log10(org_stft.abs()+eps)
+    gen_mag = torch.log10(gen_stft.abs()+eps)
 
     # temporal invariant phase
-    phase_loss = compute_phase_loss(org_temp,gen_temp)
+    phase_loss = compute_phase_loss(org_mag.sum(-2,keepdim=True),gen_mag.sum(-2,keepdim=True))
 
     # temporal envelope
-    gen_temp = compute_temporal_envelope(gen_mag, kernel_size=kernel_size)
-    org_temp = compute_temporal_envelope(org_mag, kernel_size=kernel_size)
+    gen_temp = compute_envelope(gen_mag,dim=-2,kernel_size=3)
+    org_temp = compute_envelope(org_mag,dim=-2,kernel_size=3)
     temporal_loss = F.smooth_l1_loss(gen_temp, org_temp)
 
     # spectral
-    gen_spec = compute_cepstrals(gen_mag,kernel_size=kernel_size)
-    org_spec = compute_cepstrals(org_mag,kernel_size=kernel_size)
+    gen_spec = compute_envelope(gen_mag,dim=-1,kernel_size=3)
+    org_spec = compute_envelope(org_mag,dim=-1,kernel_size=3)
     spectral_loss = F.smooth_l1_loss(gen_spec, org_spec)
 
     return temporal_loss, spectral_loss, phase_loss
 
-def compute_harmonics(stft_matrix: torch.Tensor, kernel_size=31):
+def compute_harmonics(stft_matrix: torch.Tensor, harmonic_kernel_sizes=[11,17,23], percussive_kernel_sizes=[3,7,13], eps=1e-6):
     # Calculate log-magnitude spectrograms
-    mag = torch.log1p(stft_matrix.abs())
-    
-    # Apply median filtering to separate harmonics and percussives
-    harmonic = median_pool1d(mag, kernel_size, stride=kernel_size//2+1)
-    percussive = median_pool1d(mag.transpose(-2,-1), kernel_size, stride=kernel_size//2+1).transpose(-2,-1)
-    
-    return harmonic.nan_to_num(0), percussive.nan_to_num(0)
+    mag = torch.log10(stft_matrix.abs()+eps)
+
+    harmonic_list = []
+    percussive_list = []
+
+    for kernel_size in harmonic_kernel_sizes:
+        harmonic_list.append(median_pool1d(mag, kernel_size=kernel_size, dim=-1).nan_to_num(0).view(mag.size(0),-1))
+
+    for kernel_size in percussive_kernel_sizes:
+        percussive_list.append(median_pool1d(mag, kernel_size=kernel_size, dim=-2).nan_to_num(0).view(mag.size(0),-1))
+
+    # Concatenate the results
+    harmonic = torch.cat(harmonic_list, dim=-1)
+    percussive = torch.cat(percussive_list, dim=-1)
+
+    return harmonic, percussive
 
 def combined_aux_loss(
         original_audio: torch.Tensor, generated_audio: torch.Tensor, sample_rate: int,
         c_mfcc=1., c_lfcc=1., c_hd=1., c_sts=1.,
         n_fft=1024, hop_length=320, win_length=1024,
         n_mfcc=13, n_lfcc=13, n_filter=80,
-        f_min=0., f_max=None):
+        f_min=0., f_max=None, eps=None):
 
     kernel_size = n_fft//hop_length+1
-    if kernel_size % 2 == 0: kernel_size+=1
+    if kernel_size % 2 == 0: kernel_size+=1 #enforce odd kernel size
+    if eps is None: eps = torch.finfo(original_audio.dtype).eps
 
     # Compute STFT once
-    if c_hd+c_sts>0:
+    if c_hd+c_sts+c_mfcc+c_lfcc>0:
+        
         hann_window = torch.hann_window(win_length).to(
             dtype=original_audio.dtype, device=original_audio.device
         )
@@ -129,11 +277,15 @@ def combined_aux_loss(
             return_complex=True,
             onesided=True,
             center=False)
+        
+        if c_sts+c_mfcc+c_lfcc>0:
+            org_mag = torch.log10(original_stft.abs()+eps)
+            gen_mag = torch.log10(generated_stft.abs()+eps)
     
     # Harmonic Loss
     if c_hd>0:
-        original_harmonics, original_percussives = compute_harmonics(original_stft,kernel_size)
-        generated_harmonics, generated_percussives = compute_harmonics(generated_stft,kernel_size)
+        original_harmonics, original_percussives = compute_harmonics(original_stft,eps=eps)
+        generated_harmonics, generated_percussives = compute_harmonics(generated_stft,eps=eps)
         # Define loss terms
         harmonic_loss = F.smooth_l1_loss(generated_harmonics, original_harmonics)
         harmonic_loss += F.smooth_l1_loss(generated_percussives, original_percussives)
@@ -143,37 +295,48 @@ def combined_aux_loss(
 
     # Spectral Temporal Smoothness Loss
     if c_sts>0:
-        temporal_loss, spectral_loss, phase_loss = compute_sts_loss(generated_stft,original_stft,kernel_size)
-        temporal_loss *= c_sts
-        spectral_loss *= c_sts
-        phase_loss *= c_sts
-    else: spectral_loss = temporal_loss = phase_loss = 0
+        # temporal invariant phase
+        phase_loss = compute_phase_loss(org_mag.sum(-2,keepdim=True),gen_mag.sum(-2,keepdim=True)) * c_sts
+        
+    else: phase_loss = 0
 
     # MFCC Loss
     if c_mfcc>0:
-        mfcc_transform = torchaudio.transforms.MFCC(
-            sample_rate=sample_rate, n_mfcc=n_mfcc, log_mels=True,
-            melkwargs=dict(n_fft=n_fft, hop_length=hop_length, win_length=win_length, n_mels=n_filter, f_min=f_min, f_max=f_max)
-        ).to(original_audio.device)
-        original_mfcc = mfcc_transform(original_audio)
-        generated_mfcc = mfcc_transform(generated_audio)
-        mfcc_loss = F.smooth_l1_loss(generated_mfcc, original_mfcc) * c_mfcc
-    else: mfcc_loss = 0
+        # temporal envelope
+        gen_temp = compute_envelope(gen_mag,dim=-2,kernel_size=3)
+        org_temp = compute_envelope(org_mag,dim=-2,kernel_size=3)
+        temporal_loss = F.smooth_l1_loss(gen_temp, org_temp) * c_mfcc
+
+        # mfcc_transform = torchaudio.transforms.MFCC(
+        #     sample_rate=sample_rate, n_mfcc=n_mfcc, log_mels=True,
+        #     melkwargs=dict(n_fft=n_fft, hop_length=hop_length, win_length=win_length, n_mels=n_filter, f_min=f_min, f_max=f_max)
+        # ).to(original_audio.device)
+        # original_mfcc = mfcc_transform(original_audio)
+        # generated_mfcc = mfcc_transform(generated_audio)
+        # mfcc_loss = F.smooth_l1_loss(generated_mfcc, original_mfcc) * c_mfcc
+        
+    else: temporal_loss = 0
 
     # LFCC Loss
     if c_lfcc>0:
-        lfcc_transform = torchaudio.transforms.LFCC(
-            sample_rate=sample_rate, n_lfcc=n_lfcc, n_filter=n_filter, f_min=f_min, f_max=f_max,log_lf=True,
-            speckwargs=dict(n_fft=n_fft, hop_length=hop_length, win_length=win_length)
-        ).to(original_audio.device)
-        original_lfcc = lfcc_transform(original_audio)
-        generated_lfcc = lfcc_transform(generated_audio)
-        lfcc_loss = F.smooth_l1_loss(generated_lfcc, original_lfcc) * c_lfcc
-    else: lfcc_loss = 0
+        # spectral
+        gen_spec = compute_envelope(gen_mag,dim=-1,kernel_size=3)
+        org_spec = compute_envelope(org_mag,dim=-1,kernel_size=3)
+        spectral_loss = F.smooth_l1_loss(gen_spec, org_spec) * c_lfcc
+
+        # lfcc_transform = torchaudio.transforms.LFCC(
+        #     sample_rate=sample_rate, n_lfcc=n_lfcc, n_filter=n_filter, f_min=f_min, f_max=f_max,log_lf=True,
+        #     speckwargs=dict(n_fft=n_fft, hop_length=hop_length, win_length=win_length)
+        # ).to(original_audio.device)
+        # original_lfcc = lfcc_transform(original_audio)
+        # generated_lfcc = lfcc_transform(generated_audio)
+        # lfcc_loss = F.smooth_l1_loss(generated_lfcc, original_lfcc) * c_lfcc
+    else: spectral_loss = 0
+    mfcc_loss=lfcc_loss=0
     
     return harmonic_loss, temporal_loss, spectral_loss, mfcc_loss, lfcc_loss, phase_loss
 
-def gradient_norm_loss(original_audio: torch.Tensor, generated_audio: torch.Tensor, net_d):
+def gradient_norm_loss(original_audio: torch.Tensor, generated_audio: torch.Tensor, net_d: torch.nn.Module):
     loss=0
     # Compute the gradient penalty
     # Randomly interpolate between real and generated data
@@ -186,15 +349,16 @@ def gradient_norm_loss(original_audio: torch.Tensor, generated_audio: torch.Tens
     _, disc_interpolated_output, _, _ = net_d(original_audio, interpolated)
     # Compute gradients of discriminator output w.r.t. interpolated data
     for output in disc_interpolated_output:
-        gradient = torch.autograd.grad(
+        net_d.zero_grad()
+        gradients = torch.autograd.grad(
             outputs=output,
             inputs=interpolated,
             grad_outputs=torch.ones(output.size(), device=original_audio.device),
-            create_graph=True,
+            retain_graph=True,
             only_inputs=True
-        )[0]
-        gradient = gradient.view(gradient.size(0), -1)
-        loss += torch.log(gradient.norm(2,dim=-1)).abs().mean() # force gradnorm=1
+        )
+        grad_norm = torch.stack([param.norm() for param in gradients if hasattr(param, "norm")]).norm(2)
+        loss += torch.log10(grad_norm).abs().mean() # force gradnorm=1
     return loss/len(disc_interpolated_output)
 
 def feature_loss(fmap_r, fmap_g):

lib/train/utils.py

@@ -105,6 +105,7 @@ def load_checkpoint(checkpoint_path, model, optimizer=None, load_opt=1):
 
     iteration = checkpoint_dict["iteration"]
     learning_rate = checkpoint_dict["learning_rate"]
+    kwargs = checkpoint_dict.get("kwargs", {})
     if (
         optimizer is not None and load_opt == 1
     ):  ###加载不了，如果是空的的话，重新初始化，可能还会影响lr时间表的更新，因此在train文件最外围catch
@@ -113,15 +114,11 @@ def load_checkpoint(checkpoint_path, model, optimizer=None, load_opt=1):
     #   except:
     #     traceback.print_exc()
     logger.info("Loaded checkpoint '{}' (epoch {})".format(checkpoint_path, iteration))
-    return model, optimizer, learning_rate, iteration
+    return model, optimizer, learning_rate, iteration, kwargs
 
 
-def save_checkpoint(model, optimizer, learning_rate, iteration, checkpoint_path):
-    logger.info(
-        "Saving model and optimizer state at epoch {} to {}".format(
-            iteration, checkpoint_path
-        )
-    )
+def save_checkpoint(model, optimizer, learning_rate, iteration, checkpoint_path, **kwargs):
+    logger.info(f"Saving model and optimizer state at epoch {iteration} to {checkpoint_path}")
     if hasattr(model, "module"):
         state_dict = model.module.state_dict()
     else:
@@ -132,6 +129,7 @@ def save_checkpoint(model, optimizer, learning_rate, iteration, checkpoint_path)
             "iteration": iteration,
             "optimizer": optimizer.state_dict(),
             "learning_rate": learning_rate,
+            "kwargs": kwargs
         },
         checkpoint_path,
     )