parse new arguments to mnist.py

Signed-off-by: Shashank Mittal <[email protected]>

Author: shashank-iitbhu

examples/v1beta1/hp-tuning/hyperopt-distribution.yaml

@@ -11,9 +11,9 @@ spec:
     objectiveMetricName: loss
   algorithm:
     algorithmName: random
-  parallelTrialCount: 5
-  maxTrialCount: 20
-  maxFailedTrialCount: 5
+  parallelTrialCount: 3
+  maxTrialCount: 12
+  maxFailedTrialCount: 3
   parameters:
     # - name: lr
     #   parameterType: double
@@ -39,6 +39,7 @@ spec:
       feasibleSpace:
         min: "0.1"
         max: "0.5"
+        step: "0.001"
         distribution: "logNormal"
   trialTemplate:
     primaryContainerName: training-container

examples/v1beta1/trial-images/pytorch-mnist/mnist.py

@@ -150,6 +150,20 @@ def main():
         metavar="M",
         help="SGD momentum (default: 0.5)",
     )
+    parser.add_argument(
+        "--weight-decay",
+        type=float,
+        default=0.01,
+        metavar="WD",
+        help="Weight decay for regularization (default: 0.01)",
+    )
+    parser.add_argument(
+        "--dropout-rate",
+        type=float,
+        default=0.5,
+        metavar="DR",
+        help="Dropout rate for the model (default: 0.5)",
+    )
     parser.add_argument(
         "--no-cuda", action="store_true", default=False, help="disables CUDA training"
     )