diff options
Diffstat (limited to 'utils')
| -rw-r--r-- | utils/configuration.py | 4 | ||||
| -rw-r--r-- | utils/triplet_loss.py | 2 |
2 files changed, 4 insertions, 2 deletions
diff --git a/utils/configuration.py b/utils/configuration.py index c4c4b4d..4ab1520 100644 --- a/utils/configuration.py +++ b/utils/configuration.py @@ -7,6 +7,7 @@ class SystemConfiguration(TypedDict): disable_acc: bool CUDA_VISIBLE_DEVICES: str save_dir: str + image_log_on: bool class DatasetConfiguration(TypedDict): @@ -31,6 +32,7 @@ class ModelHPConfiguration(TypedDict): ae_feature_channels: int f_a_c_p_dims: tuple[int, int, int] hpm_scales: tuple[int, ...] + hpm_use_1x1conv: bool hpm_use_avg_pool: bool hpm_use_max_pool: bool fpfe_feature_channels: int @@ -40,7 +42,7 @@ class ModelHPConfiguration(TypedDict): tfa_squeeze_ratio: int tfa_num_parts: int embedding_dims: int - triplet_margin: float + triplet_margins: tuple[float, float] class SubOptimizerHPConfiguration(TypedDict): diff --git a/utils/triplet_loss.py b/utils/triplet_loss.py index 8c143d6..d573ef4 100644 --- a/utils/triplet_loss.py +++ b/utils/triplet_loss.py @@ -34,5 +34,5 @@ class BatchAllTripletLoss(nn.Module): parted_loss_mean = all_loss.sum(1) / non_zero_counts parted_loss_mean[non_zero_counts == 0] = 0 - loss = parted_loss_mean.sum() + loss = parted_loss_mean.mean() return loss |