Add loss weights to training script

Sharda-Tech · Dec 11, 2019 · 99f1d0d · 99f1d0d
1 parent a3f9f65
commit 99f1d0d
Showing 1 changed file with 9 additions and 3 deletions.
diff --git a/examples/discovery_of_latent_keypoints/discover_latent_keypoints.py b/examples/discovery_of_latent_keypoints/discover_latent_keypoints.py
@@ -38,6 +38,10 @@
                     help='Number of filters in convolutional blocks')
 parser.add_argument('-bs', '--batch_size', default=20, type=int,
                     help='Batch size for training')
+parser.add_argument('-lw', '--loss_weights', nargs='+', type=float,
+                    default=[1.0, 1.0, 1.0, 0.2, 0.5],
+                    help='Loss weights in the following order:'
+                    '[consistency, silhouette, separation, pose, variance]')
 parser.add_argument('-lr', '--learning_rate', default=0.001, type=float,
                     help='Initial learning rate for Adam')
 parser.add_argument('-is', '--image_size', default=128, type=int,
@@ -95,13 +99,15 @@
                           args.depth * 10, args.filters, args.alpha)
 
 # loss instantiation
-loss = KeypointNetLoss(args.num_keypoints, focal_length)
+name = ['consistency', 'silhouette', 'separation', 'relative_pose', 'variance']
+weights = dict(zip(name, args.loss_weights))
+loss = KeypointNetLoss(args.num_keypoints, focal_length, loss_weights=weights)
 losses = {'uvz_points-shared': loss.uvz_points,
           'uv_volumes-shared': loss.uv_volumes}
-uvz_point_losses = [loss.consistency, loss.separation, loss.relative_pose]
 
 # metrics
-metrics = {'uvz_points-shared': uvz_point_losses,
+uvz_point_metrics = [loss.consistency, loss.separation, loss.relative_pose]
+metrics = {'uvz_points-shared': uvz_point_metrics,
            'uv_volumes-shared': [loss.silhouette, loss.variance]}
 
 # model compilation