From 282f81df38ca448eeae1695075d781d3b59e8c10 Mon Sep 17 00:00:00 2001
From: Shashank Mittal <shashank.mittal.mec22@itbhu.ac.in>
Date: Sat, 7 Sep 2024 20:17:29 +0530
Subject: [PATCH] added e2e tests for NORMAL and LOG_NORMAL

Signed-off-by: Shashank Mittal <shashank.mittal.mec22@itbhu.ac.in>

sigma calculation fixed

fix

parse new arguments to mnist.py
---
 .../hp-tuning/hyperopt-distribution.yaml      | 21 ++++++++++++++++++
 .../trial-images/pytorch-mnist/mnist.py       | 14 ++++++++++++
 .../v1beta1/hyperopt/base_service.py          | 22 ++++++++++---------
 3 files changed, 47 insertions(+), 10 deletions(-)

diff --git a/examples/v1beta1/hp-tuning/hyperopt-distribution.yaml b/examples/v1beta1/hp-tuning/hyperopt-distribution.yaml
index ea326124c27..cde00798adc 100644
--- a/examples/v1beta1/hp-tuning/hyperopt-distribution.yaml
+++ b/examples/v1beta1/hp-tuning/hyperopt-distribution.yaml
@@ -28,6 +28,19 @@ spec:
         min: "0.5"
         max: "0.9"
         distribution: "logUniform"
+    - name: weight_decay
+      parameterType: double
+      feasibleSpace:
+        min: "0.01"
+        max: "0.05"
+        distribution: "normal"
+    - name: dropout_rate
+      parameterType: double
+      feasibleSpace:
+        min: "0.1"
+        max: "0.5"
+        step: "0.001"
+        distribution: "logNormal"
   trialTemplate:
     primaryContainerName: training-container
     trialParameters:
@@ -37,6 +50,12 @@ spec:
       - name: momentum
         description: Momentum for the training model
         reference: momentum
+      - name: weightDecay
+        description: Weight decay for the training model
+        reference: weight_decay
+      - name: dropoutRate
+        description: Dropout rate for the training model
+        reference: dropout_rate
     trialSpec:
       apiVersion: batch/v1
       kind: Job
@@ -53,6 +72,8 @@ spec:
                   - "--batch-size=16"
                   - "--lr=${trialParameters.learningRate}"
                   - "--momentum=${trialParameters.momentum}"
+                  - "--weight-decay=${trialParameters.weightDecay}"
+                  - "--dropout-rate=${trialParameters.dropoutRate}"
                 resources:
                   limits:
                     memory: "1Gi"
diff --git a/examples/v1beta1/trial-images/pytorch-mnist/mnist.py b/examples/v1beta1/trial-images/pytorch-mnist/mnist.py
index 7ecc911cbb4..91b46e31449 100644
--- a/examples/v1beta1/trial-images/pytorch-mnist/mnist.py
+++ b/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"
     )
diff --git a/pkg/suggestion/v1beta1/hyperopt/base_service.py b/pkg/suggestion/v1beta1/hyperopt/base_service.py
index 0d65ecd77e0..283942e8c3b 100644
--- a/pkg/suggestion/v1beta1/hyperopt/base_service.py
+++ b/pkg/suggestion/v1beta1/hyperopt/base_service.py
@@ -93,34 +93,36 @@ def create_hyperopt_domain(self):
                             param.name, float(param.min), float(param.max)
                         )
                 elif param.distribution == api_pb2.NORMAL:
-                    sigma = 1
+                    mu = (float(param.min) + float(param.max)) / 2
+                    sigma = (float(param.max) - float(param.min)) / 6
                     if param.step:
                         hyperopt_search_space[param.name] = hyperopt.hp.qnormal(
                             param.name,
-                            float((float(param.min) + float(param.max)) / 2),
-                            float(sigma),
+                            mu,
+                            sigma,
                             float(param.step),
                         )
                     else:
                         hyperopt_search_space[param.name] = hyperopt.hp.normal(
                             param.name,
-                            float((float(param.min) + float(param.max)) / 2),
-                            float(sigma),
+                            mu,
+                            sigma,
                         )
                 elif param.distribution == api_pb2.LOG_NORMAL:
-                    sigma = 1
+                    mu = (float(param.min) + float(param.max)) / 2
+                    sigma = (float(param.max) - float(param.min)) / 6
                     if param.step:
                         hyperopt_search_space[param.name] = hyperopt.hp.qlognormal(
                             param.name,
-                            float((float(param.min) + float(param.max)) / 2),
-                            float(sigma),
+                            mu,
+                            sigma,
                             float(param.step),
                         )
                     else:
                         hyperopt_search_space[param.name] = hyperopt.hp.lognormal(
                             param.name,
-                            float((float(param.min) + float(param.max)) / 2),
-                            float(sigma),
+                            mu,
+                            sigma,
                         )
             elif param.type == CATEGORICAL or param.type == DISCRETE:
                 hyperopt_search_space[param.name] = hyperopt.hp.choice(