Update Readme and Jax_Utils

README.md

@@ -60,20 +60,22 @@ Please extract the zip file and replace the original ``./data/`` folder by the e
 - Jaxlib == 0.3.25+cuda11.cudnn82
 - jaxopt == 0.5.5
 - optax == 0.1.4
+- chex == 0.1.5
 - dm-haiku == 0.0.9
 
-To install the dependencies: `pip3 install -r requirements.txt`
+Please follow the installation guidelines in [http://github.com/google/jax](https://github.com/google/jax#pip-installation-gpu-cuda) to install compatible version Jax and Jaxlib version for your machine. The version of packages related to Jax (Jaxopt, Optax, Chex) may also need to be adjusted for compatible with Jax and Jaxlib.
 
+To install other dependencies: `pip3 install -r requirements.txt`
 
 ## Experiments
 
 The general template commands for running an experiment are:
 
 ```bash
-bash runners/<dataset>/run_<algorithm>.sh --repeat 1 [other flags]
+bash runners/<dataset>/run_<algorithm>.sh [other flags]
 ```
 
-### Basic Usage
+### Flags
 
 | params | full params        | description                                               | default value | options |
 |--------|--------------------|-----------------------------------------------------------|---------------|---------|
@@ -109,10 +111,10 @@ bash runners/shakespeare/run_fedeq.sh -r 1 -g 0
 Run algorithms to evaluate the generalization to unseen clients
 
 ```bash
-bash runners/femnist/run_fedeq.sh -r 1 -fu 0.1 -g 0
-bash runners/cifar10/run_fedeq.sh -r 1 -fu 0.1 -g 0
-bash runners/cifar100/run_fedeq.sh -r 1 -fu 0.1 -g 0
-bash runners/shakespeare/run_fedeq.sh -r 1 -fu 0.1 -g 0
+bash runners/femnist/run_fedeq.sh -fu 0.1 -r 1 -g 0
+bash runners/cifar10/run_fedeq.sh -fu 0.1 -r 1 -g 0
+bash runners/cifar100/run_fedeq.sh -fu 0.1 -r 1 -g 0
+bash runners/shakespeare/run_fedeq.sh -fu 0.1 -r 1 -g 0
 ```
 
 ---

requirements.txt

@@ -1,9 +1,13 @@
+jaxopt == 0.5.5
+dm-haiku == 0.0.9
+optax == 0.1.4
+chex == 0.1.5
 numpy
 tqdm
 scikit-learn
 pandas
-jax==0.3.25
-jaxlib==0.3.25+cuda11.cudnn82
-jaxopt==0.5.5
-dm-haiku==0.0.9
-optax==0.1.4
+torch
+torchvision
+seaborn
+matplotlib
+tqdm

runners/shakespeare/run_fedeq.sh

@@ -1,3 +1,3 @@
 bash runners/shakespeare/main.sh \
---trainer fedeq_seq --model deq_transformer --rho 0.01 --lam_admm 0.01 \
+--trainer fedeq_seq --model deq_transformer --rho 0.01 --lam_admm 0.001 \
 $@

utils/jax_utils.py

@@ -211,7 +211,7 @@ def rdeq_nll_loss_fn(model,
   """Computes the (scalar) LM loss on `data` w.r.t. params."""
   inputs, targets = batch
 
-  params = hk.data_structures.merge(shared_params, personalized_params)
+  params = hk.data_structures.merge(deq_params, personalized_params)
   logits, embedding = model.apply(params, rng, inputs)
   targets = jax.nn.one_hot(targets, vocab_size)
   assert logits.shape == targets.shape
@@ -222,6 +222,7 @@ def rdeq_nll_loss_fn(model,
   quadratic_term = 0.5 * rho * tree_l2_norm(diff_params, squared=True)
 
   mask = jnp.greater(inputs, 0)
+  #print(jnp.sum(mask))
   log_likelihood = jnp.sum(targets * jax.nn.log_softmax(logits), axis=-1)
   return -jnp.sum(log_likelihood * mask)/jnp.sum(mask) + linear_term/jnp.sum(mask)  + quadratic_term/jnp.sum(mask)  # NLL per token.