new optimizer for 4f system

Author: CarlaRodriguez96

experiments/four_f_optimizer.py

@@ -17,64 +17,91 @@
 # Call the data loader and set batchsize
 dataloader = MultiHDF5DataLoader("training_data_4f", batch_size=10)
 
-# Define the update for batch optimization:
-@jit
-def update(step_index, optimizer_state, input_fields, target_fields):
-    parameters = get_params(optimizer_state)
-    # Call the loss function and compute the gradients
-    computed_loss = loss_fn(parameters, input_fields, target_fields)
-    computed_gradients = grad(loss_fn, allow_int=True)(parameters, input_fields, target_fields)
-    return opt_update(step_index, computed_gradients, optimizer_state), computed_loss, computed_gradients
+# Define the loss function and compute its gradients:
+loss_function = jit(loss_dualSLM)
 
-# JIT the loss function:
-loss_fn = jit(loss_dualSLM)
+# ----------------------------------------------------
 
-# Optimizer settings
-STEP_SIZE = 0.1
-num_iterations = 50000
-n_best = 500
-best_loss = 1e2
-best_params = None
-best_step = 0
-
-# Init random parameters
-phase_mask_slm1 = jnp.array([np.random.uniform(0, 1, (shape, shape))], dtype=jnp.float64)[0]
-phase_mask_slm2 = jnp.array([np.random.uniform(0, 1, (shape, shape))], dtype=jnp.float64)[0]
-distance_0 = jnp.array([np.random.uniform(0.027, 1)], dtype=jnp.float64)
-distance_1 = jnp.array([np.random.uniform(0.027, 1)], dtype=jnp.float64) 
-distance_2 = jnp.array([np.random.uniform(0.027, 1)], dtype=jnp.float64)
-init_params = [distance_0, distance_1, distance_2, phase_mask_slm1, phase_mask_slm2]
+def fit(params: optax.Params, optimizer: optax.GradientTransformation, num_iterations) -> optax.Params:
 
-# Define the optimizer and initialize it
-opt_init, opt_update, get_params = optimizers.adam(STEP_SIZE)
-opt_state = opt_init(init_params)
-        
-# Optimize in a loop:
-print('Starting Optimization', flush=True)    
-tic = time.perf_counter()
+    opt_state = optimizer.init(params)
+    
+    @jit
+    def update(params, opt_state, input_fields, target_fields):
+        # Define single update step:
+        # JIT the loss and compute 
+        loss_value, grads = jax.value_and_grad(loss_function, allow_int=True)(params, input_fields, target_fields)
+        # Update the state of the optimizer
+        updates, opt_state = optimizer.update(grads, opt_state, params)
+        params = optax.apply_updates(params, updates)
+        return params, opt_state, loss_value
 
-for step in range(num_iterations):
+    # Initialize some parameters
+    iteration_steps=[]
+    loss_list=[]
+
+    # Optimizer settings
+    n_best = 500
+    best_loss = 1e2
+    best_params = None
+    best_step = 0
 
-    # Load data:    
-    input_fields, target_fields = next(dataloader)
-    # Perform an update step:
-    opt_state, loss_value, gradients = update(step, opt_state, input_fields, target_fields)
+    print('Starting Optimization', flush=True)
 
-    # Update the `best_loss` value:
-    if loss_value < best_loss:
-        # Best loss value
-        best_loss = loss_value
-        # Best optimized parameters
-        best_params = get_params(opt_state)
-        best_step = step
-        print('Best loss value is updated')
+    for step in range(num_iterations):        
+        # Load data:
+        input_fields, target_fields = next(dataloader)
+        params, opt_state, loss_value = update(params, opt_state, input_fields, target_fields)
+        
+        print(f"Step {step}")
+        print(f"Loss {loss_value}")
+        
+        iteration_steps.append(step)
+        loss_list.append(loss_value)
+        
+        # Update the `best_loss` value:
+        if loss_value < best_loss:
+            # Best loss value
+            best_loss = loss_value
+            # Best optimized parameters
+            best_params = params
+            best_step = step
+            print('Best loss value is updated')
+            
+        if step % 100 == 0:
+            # Stopping criteria: if best_loss has not changed every 500 steps, stop.
+            if step - best_step > n_best:
+                print(f'Stopping criterion: no improvement in loss value for {n_best} steps')
+                break
+    
+    print(f'Best loss: {best_loss} at step {best_step}')
+    print(f'Best parameters: {best_params}')  
+    return best_params, best_loss, iteration_steps, loss_list
+
+# ----------------------------------------------------
 
-    if step % 500 == 0:
-        # Stopping criteria: if best_loss has not changed every 500 steps, stop.
-        if step - best_step > n_best:
-            print(f'Stopping criterion: no improvement in loss value for {n_best} steps')
-            break
+# Optimizer settings
+num_iterations = 50000
+num_samples = 50
+# Step size engineering:
+STEP_SIZE = 0.01
+WEIGHT_DECAY = 0.0001
 
-print(f'Best loss: {best_loss} at step {best_step}')
-print(f'Best parameters: {best_params}')
-print("Time taken to optimize - in seconds", time.perf_counter() - tic)
+for i in range(num_samples):
+    tic = time.perf_counter()
+    
+    # Init random parameters
+    phase_mask_slm1 = jnp.array([np.random.uniform(0, 1, (shape, shape))], dtype=jnp.float64)[0]
+    phase_mask_slm2 = jnp.array([np.random.uniform(0, 1, (shape, shape))], dtype=jnp.float64)[0]
+    distance_0 = jnp.array([np.random.uniform(0.027, 1)], dtype=jnp.float64)
+    distance_1 = jnp.array([np.random.uniform(0.027, 1)], dtype=jnp.float64) 
+    distance_2 = jnp.array([np.random.uniform(0.027, 1)], dtype=jnp.float64)
+    init_params = [distance_0, distance_1, distance_2, phase_mask_slm1, phase_mask_slm2]
+    
+    # Init optimizer:
+    optimizer = optax.adamw(STEP_SIZE, weight_decay=WEIGHT_DECAY)
+    
+    # Apply fit function:
+    best_params, best_loss, iteration_steps, loss_list = fit(init_params, optimizer, num_iterations)
+    
+    print("Time taken to optimize one sample - in seconds", time.perf_counter() - tic)