Upload 10 files

README.md

@@ -1,3 +1,185 @@
----
-license: unknown
----
+<h1 align="center">
+  <img src="images/ProARD_logo.png" width="500"/>
+  <br/>
+    PROARD: PROGRESSIVE ADVERSARIAL ROBUSTNESS DISTILLATION: PROVIDE WIDE RANGE OF ROBUST STUDENTS
+  </br>
+</h1>
+<p align="center">
+<a href="#background">Background</a> •
+<a href="#usage">Usage</a> •
+<a href="#code">Code</a> •
+<a href="#citation">Citation</a> •
+</p>
+
+## Background
+Progressive Adversarial Robustness Distillation (ProARD), enabling the efficient
+one-time training of a dynamic network that supports a diverse range of accurate and robust student
+networks without requiring retraining. ProARD makes a dynamic deep neural network based on
+dynamic layers by encompassing variations in width, depth, and expansion in each design stage to
+support a wide range of architectures. 
+
+<h1 align="center">
+  <img src="images/ProARD.png" width="1000"/>
+</h1>
+
+
+
+## Usage
+```
+git clone https://github.com/hamidmousavi0/ProARD.git
+```
+## Code Structure
+```
+- attacks/ # Different Adversarial attack methods (PGD, AutoAttack, FGSM, DeepFool, etc. ([Refrence](https://github.com/imrahulr/hat.git)))
+- proard/ 
+    - classification/ 
+        - data_provider/ # The dataset and dataloader definitions for Cifar-10, Cifar-100, and ImageNet.
+        - elastic_nn/ 
+            - modules/ # The deficnition of dynamic layers 
+            - networks/ # The deficnition of dynamic networks 
+            - training/ # Progressive training 
+        -networks/ # The original networks
+        -run_anager/ # The Configs and distributed training   
+    - nas 
+        - accuracy_predictor/ # The accuracy and robustness predictor
+        - efficiency_predictor/ # The efficiency predictor
+        - search_algorithm/ # The Multi-Objective Search Engine
+    - utils/ # Utility functions
+    - model_zoo.py # All the models for evaluation   
+- create_acc_rob_pred_dataset.py # Create dataset to train the accuracy-robustness predictor.
+- create_acc_rob_pred.py # make the predictor model.
+- eval_ofa_net.py # Eval the sub-nets
+- search_best.py # Search the best sub-net
+- train_ofa_net_WPS.py # train the dynamic network without progressive training.
+- train_ofa_net.py # Train the dynamic network with progressive training.
+- train_teacher_net.py # Train teacher network for Robust knoweldge distillation. 
+
+```
+### Installing
+
+**From Source**
+
+Download this repository into your project folder.
+
+### Details of the usage
+
+## Evaluation
+
+```
+python eval_ofa_net.py --path path of dataset  --net Dynamic net name (ResNet50, MBV3) 
+                       --dataset (cifar10, cifar100) --robust_mode (True, False)
+                       --WPS (True, False) --attack ('fgsm', 'linf-pgd', 'fgm', 'l2-pgd', 'linf-df', 'l2-df', 'linf-apgd', 'l2-apgd','squar_attack','autoattack','apgd_ce')
+```
+
+
+## Training
+
+### Step-0: Train Teacher Net
+
+```
+horovodrun -np 4 python train_teacher_net.py --model_name ("ResNet50", "MBV3") --dataset (cifar10, cifar100)
+                                             --robust_mode (True, False) --epsilon 0.031 --num_steps 10
+                                             --step_size 0.0078 --distance 'l-inf' --train_criterion 'trades'
+                                             --attack_type 'linf-pgd'
+```
+
+### Step-1: Dynamic Width/Kernel training
+
+
+```
+horovodrun -np 4 python train_ofa_net.py --task 'width' or 'kernel' --model_name ("ResNet50", "MBV3") --dataset (cifar10, cifar100)
+                                             --robust_mode (True, False) --epsilon 0.031 --num_steps 10
+                                             --step_size 0.0078 --distance 'l-inf' --train_criterion 'trades'
+                                             --attack_type 'linf-pgd' --kd_criterion 'rslad' --phase 1
+```
+
+### Step-2: Dynamic Width/Kernel and depth training
+
+##### Phase-1
+```
+horovodrun -np 4 python train_ofa_net.py --task 'depth' --model_name ("ResNet50", "MBV3") --dataset (cifar10, cifar100)
+                                             --robust_mode (True, False) --epsilon 0.031 --num_steps 10
+                                             --step_size 0.0078 --distance 'l-inf' --train_criterion 'trades'
+                                             --attack_type 'linf-pgd' --kd_criterion 'rslad' --phase 1
+```
+##### Phase-2
+```
+horovodrun -np 4 python train_ofa_net.py --task 'depth' --model_name ("ResNet50", "MBV3") --dataset (cifar10, cifar100)
+                                             --robust_mode (True, False) --epsilon 0.031 --num_steps 10
+                                             --step_size 0.0078 --distance 'l-inf' --train_criterion 'trades'
+                                             --attack_type 'linf-pgd' --kd_criterion 'rslad' --phase 2
+```
+
+### Step-3: Dynamic Width/Kernel, depth, and expand training
+
+
+##### Phase-1
+```
+horovodrun -np 4 python train_ofa_net.py --task 'expand' --model_name ("ResNet50", "MBV3") --dataset (cifar10, cifar100)
+                                             --robust_mode (True, False) --epsilon 0.031 --num_steps 10
+                                             --step_size 0.0078 --distance 'l-inf' --train_criterion 'trades'
+                                             --attack_type 'linf-pgd' --kd_criterion 'rslad' --phase 1
+```
+##### Phase-2
+```
+horovodrun -np 4 python train_ofa_net.py --task 'expand' --model_name ("ResNet50", "MBV3") --dataset (cifar10, cifar100)
+                                             --robust_mode (True, False) --epsilon 0.031 --num_steps 10
+                                             --step_size 0.0078 --distance 'l-inf' --train_criterion 'trades'
+                                             --attack_type 'linf-pgd' --kd_criterion 'rslad' --phase 2
+```
+
+
+
+
+
+<!--
+
+* **ProAct** (the proposed algorithm) ([paper](https://arxiv.org/abs/2406.06313) and ([code](https://github.com/hamidmousavi0/reliable-relu-toolbox/tree/master/rrelu/search_bound/proact.py)).
+* **FitAct** ([paper](https://arxiv.org/pdf/2112.13544) and [code](https://github.com/hamidmousavi0/reliable-relu-toolbox/tree/master/rrelu/search_bound/fitact.py)).
+* **FtClipAct** ([paper](https://arxiv.org/pdf/1912.00941) and [code](https://github.com/hamidmousavi0/reliable-relu-toolbox/tree/master/rrelu/search_bound/ftclip.py)).
+* **Ranger** ([paper](https://arxiv.org/pdf/2003.13874) and [code](https://github.com/hamidmousavi0/reliable-relu-toolbox/tree/master/rrelu/search_bound/ranger.py)).
+-->
+
+
+
+
+<!-- Use the following notebook to learn the main steps of the tool.
+[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://github.com/hamidmousavi0/reliable-relu-toolbox/blob/master/RReLU.ipynb)-->
+
+
+## To-do list
+- [ ] Add object detection Task
+- [ ] Add Transformers architectures
+
+<!--
+### Run search in the command line 
+
+When you download this repository into your project folder.
+```
+torchrun --nproc_per_node=2  search.py --dataset cifar10 --data_path "./dataset/cifar10"  --batch_size 128 --model "resnet20" --n_worker 32 \
+                      --name_relu_bound "zero"  --name_serach_bound "ranger" --bounds_type "layer" --bitflip "fixed" --image_size 32 --pretrained_model
+```
+-->
+## Citation
+
+View the [published paper(preprint), Accepted in IJCNN 2025](https://www.arxiv.org/pdf/2506.07666).
+<!--
+```
+@article{mousavi2024proact,
+  title={ProAct: Progressive Training for Hybrid Clipped Activation Function to Enhance Resilience of DNNs},
+  author={Mousavi, Seyedhamidreza and Ahmadilivani, Mohammad Hasan and Raik, Jaan and Jenihhin, Maksim and Daneshtalab, Masoud},
+  journal={arXiv preprint arXiv:2406.06313},
+  year={2024}
+}
+```
+-->
+## Acknowledgment
+
+We acknowledge the National Academic Infrastructure for Supercomputing in Sweden (NAISS), partially funded by the Swedish Research Council through grant agreement no
+
+## Contributors
+Some of the code in this repository is based on the following amazing works:
+
+[Once-For-All](https://github.com/mit-han-lab/once-for-all.git)
+
+[Hat](https://github.com/imrahulr/hat.git)

create_acc_rob_pred.py

@@ -0,0 +1,176 @@
+import os
+import torch
+import argparse
+import torch.nn as nn 
+from tqdm.auto import tqdm
+from torch.utils.data import DataLoader
+import torch.nn as nn 
+import torch
+from torch import nn
+from torch.optim import *
+from torch.optim.lr_scheduler import *
+from torch.utils.data import DataLoader
+from torchprofile import profile_macs
+from torchvision.datasets import *
+from torchvision.transforms import *
+from proard.classification.data_providers.imagenet import ImagenetDataProvider
+from proard.classification.run_manager import DistributedClassificationRunConfig, DistributedRunManager
+from proard.model_zoo import DYN_net
+from proard.nas.accuracy_predictor import AccuracyDataset,AccuracyPredictor,ResNetArchEncoder,RobustnessPredictor,MobileNetArchEncoder,AccuracyRobustnessDataset,Accuracy_Robustness_Predictor
+parser = argparse.ArgumentParser()
+
+
+def RMSELoss(yhat,y):
+    return torch.sqrt(torch.mean((yhat-y)**2))
+def train(
+  model: nn.Module,
+  dataloader: DataLoader,
+  criterion: nn.Module,
+  optimizer: Optimizer,
+  callbacks = None,
+  epochs = 10,
+  save_path = None
+) -> None:
+  model.cuda()
+  model.train()
+  for epoch in range(epochs):
+    print(epoch)
+    for inputs, targets_acc, targets_rob in tqdm(dataloader, desc='train', leave=False):
+      inputs = inputs.float().cuda()
+      targets_acc = targets_acc.cuda()
+      targets_rob = targets_rob.cuda()
+
+      # Reset the gradients (from the last iteration)
+      optimizer.zero_grad()
+
+      # Forward inference
+      outputs = model(inputs)
+      loss = criterion(outputs[:,0], targets_acc) + criterion(outputs[:,1], targets_rob)
+
+      # Backward propagation
+      loss.backward()
+
+      # Update optimizer and LR scheduler
+      optimizer.step()
+      # scheduler.step(epoch)
+
+      if callbacks is not None:
+          for callback in callbacks:
+              callback()        
+  torch.save(model.state_dict(), save_path)
+  return model
+
+@torch.inference_mode()
+def evaluate(
+  model: nn.Module,
+  dataloader: DataLoader,
+) -> float:
+  model.eval()
+
+  for inputs, targets_acc, targets_rob in tqdm(dataloader, desc="eval", leave=False):
+    # Move the data from CPU to GPU
+    inputs = inputs.cuda()
+
+    targets_acc = targets_acc.cuda()
+    targets_rob = targets_rob.cuda()
+
+
+    # Inference
+    outputs = model(inputs)
+
+    # Convert logits to class indices
+    print(RMSELoss(outputs[:,0],targets_acc),RMSELoss(outputs[:,1],targets_rob))
+  return RMSELoss(outputs[:,0],targets_acc) + RMSELoss(outputs[:,1],targets_rob)
+
+
+def get_model_flops(model, inputs):
+    num_macs = profile_macs(model, inputs)
+    return num_macs
+
+
+def get_model_size(model: nn.Module, data_width=32):
+    """
+    calculate the model size in bits
+    :param data_width: #bits per element
+    """
+    num_elements = 0
+    for param in model.parameters():
+        num_elements += param.numel()
+    return num_elements * data_width
+
+
+
+
+
+parser.add_argument(
+    "-p", "--path", help="The path of cifar10", type=str, default="/dataset/cifar10"
+)
+parser.add_argument("-g", "--gpu", help="The gpu(s) to use", type=str, default="all")
+parser.add_argument(
+    "-b",
+    "--batch_size",
+    help="The batch on every device for validation",
+    type=int,
+    default=32,
+)
+parser.add_argument("-j", "--workers", help="Number of workers", type=int, default=20)
+parser.add_argument(
+    "-n",
+    "--net",
+    metavar="DYNNET",
+    default="ResNet50",
+    choices=[
+         "ResNet50",
+        "MBV3",
+        "ProxylessNASNet",
+    ],
+    help="Dyanmic networks",
+)
+parser.add_argument(
+    "--dataset", type=str, default="cifar10" ,choices=["cifar10", "cifar100", "imagenet"]
+)
+parser.add_argument("--train_criterion", type=str, default="trades",choices=["trades","sat","mart","hat"])
+parser.add_argument(
+    "--robust_mode", type=bool, default=True
+)
+args = parser.parse_args()
+if args.net == "ResNet50":   
+   arch = ResNetArchEncoder(image_size_list=[224 if args.dataset == 'imagenet' else 32],depth_list=[0,1,2],expand_list=[0.2,0.25,0.35],width_mult_list=[0.65,0.8,1.0])
+else:
+   arch =  MobileNetArchEncoder (image_size_list=[224 if args.dataset == 'imagenet' else 32],depth_list=[2,3,4],expand_list=[3,4,6],ks_list=[3,5,7])
+print(arch)
+acc_data = AccuracyRobustnessDataset("./acc_rob_data_{}_{}_{}".format(args.dataset,args.net,args.train_criterion))
+train_loader, valid_loader, base_acc ,base_rob = acc_data.build_acc_data_loader(arch)
+acc_pred_network = Accuracy_Robustness_Predictor(arch_encoder=arch,base_acc_val=None)
+# optimizer_ = torch.optim.Adam(acc_pred_network.parameters(),lr=1e-3,weight_decay=1e-4)
+# criterion = nn.MSELoss()
+# acc_pred_network = train(acc_pred_network,train_loader,criterion,optimizer_,callbacks=None, epochs=50,save_path ="./acc_rob_data_{}_{}_{}/src/model_acc_rob.pth".format(args.dataset,args.net,args.train_criterion).format(args.dataset))
+acc_pred_network.load_state_dict(torch.load("./acc_rob_data_{}_{}_{}/src/model_acc_rob.pth".format(args.dataset,args.net,args.train_criterion)))
+print(evaluate(acc_pred_network,valid_loader))
+
+
+
+# import numpy as np 
+# accs=[]
+# robs=[]
+# pred_accs=[]
+# pred_robs=[]
+# for x,acc,rob, in valid_loader:
+#    for ac in acc:
+#       accs.append(ac.item()*100)
+#    for ro in rob:
+#       robs.append(ro.item()*100)  
+
+# for x,acc,rob, in valid_loader:
+#    for arch in x:
+#       acc ,rob = acc_pred_network(arch.cuda())
+#       pred_accs.append(acc.item()*100)
+#       pred_robs.append(rob.item()*100)
+# print(accs,robs)      
+# print(pred_accs,pred_robs)   
+# np.savetxt("./results/accs.csv", np.array(accs), delimiter=",")  
+# np.savetxt("./results/robs.csv", np.array(robs), delimiter=",") 
+# np.savetxt("./results/pred_accs.csv", np.array(pred_accs), delimiter=",")  
+# np.savetxt("./results/pred_robs.csv", np.array(pred_robs), delimiter=",")
+
+

create_acc_rob_pred_dataset.py

@@ -0,0 +1,75 @@
+import os
+import torch
+import argparse
+
+from proard.classification.data_providers.imagenet import ImagenetDataProvider
+from proard.classification.run_manager import DistributedClassificationRunConfig, DistributedRunManager
+from proard.model_zoo import DYN_net
+from proard.nas.accuracy_predictor import AccuracyRobustnessDataset
+import horovod.torch as hvd
+parser = argparse.ArgumentParser()
+parser.add_argument(
+    "-p", "--path", help="The path of cifar10", type=str, default="/dataset/cifar10"
+)
+parser.add_argument("-g", "--gpu", help="The gpu(s) to use", type=str, default="all")
+parser.add_argument(
+    "-b",
+    "--batch_size",
+    help="The batch on every device for validation",
+    type=int,
+    default=32,
+)
+parser.add_argument("-j", "--workers", help="Number of workers", type=int, default=20)
+parser.add_argument(
+    "-n",
+    "--net",
+    metavar="DYNNET",
+    default="ResNet50",
+    choices=[
+        "ResNet50",
+        "MBV3",
+        "ProxylessNASNet",
+        "MBV2"
+    ],
+    help="Dynamic networks",
+)
+parser.add_argument(
+    "--dataset", type=str, default="cifar10" ,choices=["cifar10", "cifar100", "imagenet"]
+)
+parser.add_argument("--train_criterion", type=str, default="trades",choices=["trades","sat","mart","hat"])
+parser.add_argument(
+    "--robust_mode", type=bool, default=True
+)
+parser.add_argument(
+    "--WPS", type=bool, default=True
+)
+parser.add_argument(
+    "--base", type=bool, default=False
+)
+# Initialize Horovod
+hvd.init()
+    # Pin GPU to be used to process local rank (one GPU per process)
+torch.cuda.set_device(hvd.local_rank())
+num_gpus = hvd.size()
+
+args = parser.parse_args()
+if args.gpu == "all":
+    device_list = range(torch.cuda.device_count())
+    args.gpu = ",".join(str(_) for _ in device_list)
+else:
+    device_list = [int(_) for _ in args.gpu.split(",")]
+os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
+args.test_batch_size = args.batch_size # * max(len(device_list), 1)
+ImagenetDataProvider.DEFAULT_PATH = args.path
+
+
+distributed_run_config = DistributedClassificationRunConfig(**args.__dict__, num_replicas=num_gpus, rank=hvd.rank())
+dyn_network = DYN_net(args.net, args.robust_mode , args.dataset, args.train_criterion, pretrained=True,run_config=distributed_run_config,WPS=args.WPS)
+compression = hvd.Compression.none
+distributed_run_manager = DistributedRunManager(".tmp/eval_subnet", dyn_network, distributed_run_config,compression,is_root=(hvd.rank() == 0),init=False)
+distributed_run_manager.save_config()
+    # hvd broadcast
+distributed_run_manager.broadcast()
+acc_data = AccuracyRobustnessDataset("./acc_rob_data_WPS_{}_{}_{}".format(args.dataset,args.net,args.train_criterion))
+
+acc_data.build_acc_rob_dataset(distributed_run_manager,dyn_network,image_size_list=[224 if args.dataset == "imagenet" else 32])

eval_ofa_net.py

@@ -0,0 +1,94 @@
+# Once for All: Train One Network and Specialize it for Efficient Deployment
+# Han Cai, Chuang Gan, Tianzhe Wang, Zhekai Zhang, Song Han
+# International Conference on Learning Representations (ICLR), 2020.
+
+import os
+import torch
+import argparse
+
+from proard.classification.data_providers.imagenet import ImagenetDataProvider
+from proard.classification.data_providers.cifar10 import Cifar10DataProvider
+from proard.classification.data_providers.cifar100 import Cifar100DataProvider
+from proard.classification.run_manager import ClassificationRunConfig, RunManager
+from proard.model_zoo import DYN_net
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument(
+    "-p", "--path", help="The path of imagenet", type=str, default="/dataset/imagenet"
+)
+parser.add_argument("-g", "--gpu", help="The gpu(s) to use", type=str, default="all")
+parser.add_argument(
+    "-b",
+    "--batch-size",
+    help="The batch on every device for validation",
+    type=int,
+    default=16,
+)
+parser.add_argument("-j", "--workers", help="Number of workers", type=int, default=20)
+parser.add_argument(
+    "-n",
+    "--net",
+    metavar="DYNET",
+    default="ResNet50",
+    choices=[
+        "ResNet50",
+        "MBV3",
+        "ProxylessNASNet",
+        "MBV2",
+        "WideResNet"
+    ],
+    help="dynamic networks",
+)
+parser.add_argument(
+    "--dataset", type=str, default="cifar10" ,choices=["cifar10", "cifar100", "imagenet"]
+)
+parser.add_argument(
+    "--attack", type=str, default="autoattack" ,choices=['fgsm', 'linf-pgd', 'fgm', 'l2-pgd', 'linf-df', 'l2-df', 'linf-apgd', 'l2-apgd','squar_attack','autoattack','apgd_ce']
+)
+parser.add_argument("--train_criterion", type=str, default="trades",choices=["trades","sat","mart","hat"])
+parser.add_argument(
+    "--robust_mode", type=bool, default=True
+)
+parser.add_argument(
+    "--WPS", type=bool, default=False
+)
+parser.add_argument(
+    "--base", type=bool, default=False
+)
+args = parser.parse_args()
+if args.gpu == "all":
+    device_list = range(torch.cuda.device_count())
+    args.gpu = ",".join(str(_) for _ in device_list)
+else:
+    device_list = [int(_) for _ in args.gpu.split(",")]
+os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
+args.batch_size = args.batch_size * max(len(device_list), 1)
+ImagenetDataProvider.DEFAULT_PATH = args.path
+
+run_config = ClassificationRunConfig(attack_type=args.attack,dataset= args.dataset, test_batch_size=args.batch_size, n_worker=args.workers,robust_mode=args.robust_mode)
+dyn_network = DYN_net(args.net,args.robust_mode,args.dataset, args.train_criterion ,pretrained=True,run_config=run_config,WPS=args.WPS,base=args.base)
+""" Randomly sample a sub-network, 
+    you can also manually set the sub-network using: 
+        dyn_network.set_active_subnet(ks=7, e=6, d=4) 
+"""
+if not args.base:
+    # dyn_network.set_active_subnet(ks=3, e=4, d=2)
+    dyn_network.set_active_subnet(d=2,e=0.35,w=1.0) 
+    # dyn_network.sample_active_subnet()
+    # dyn_network.set_max_net()
+    subnet = dyn_network.get_active_subnet(preserve_weight=True)
+    # print(subnet)
+else:
+    subnet = dyn_network    
+""" Test sampled subnet 
+"""
+run_manager = RunManager(".tmp/eval_subnet", subnet, run_config, init=False)
+run_config.data_provider.assign_active_img_size(32)
+run_manager.reset_running_statistics(net=subnet)
+
+print("Test random subnet:")
+# print(subnet.module_str)
+
+loss, (top1, top5,robust1,robust5) = run_manager.validate(net=subnet,is_test=True)
+print("Results: loss=%.5f,\t top1=%.1f,\t top5=%.1f,\t robust1=%.1f,\t robust5=%.1f" % (loss, top1, top5,robust1,robust5))

hugging_face.py

@@ -0,0 +1,21 @@
+from huggingface_hub import interpreter_login
+from huggingface_hub import upload_folder, delete_folder, upload_file
+# interpreter_login()
+
+
+
+# upload_folder(folder_path = "attacks/",path_in_repo="attacks", repo_id="smi08/ProArd")
+# upload_folder(folder_path = "images/",path_in_repo="images", repo_id="smi08/ProArd")
+# upload_folder(folder_path = "proard/",path_in_repo="proard", repo_id="smi08/ProArd")
+# upload_folder(folder_path = "robust_loss/",path_in_repo="robust_loss", repo_id="smi08/ProArd")
+# upload_folder(folder_path = "utils/",path_in_repo="utils", repo_id="smi08/ProArd")
+# delete_folder(path_in_repo="smi08", repo_id="smi08/ProArd")
+upload_file(path_or_fileobj="create_acc_rob_pred_dataset.py",path_in_repo="",repo_id="smi08/ProArd")
+upload_file(path_or_fileobj="create_acc_rob_pred.py",path_in_repo="",repo_id="smi08/ProArd")
+upload_file(path_or_fileobj="eval_ofa_net.py",path_in_repo="",repo_id="smi08/ProArd")
+upload_file(path_or_fileobj="sample_eval.py",path_in_repo="",repo_id="smi08/ProArd")
+upload_file(path_or_fileobj="search_best.py",path_in_repo="",repo_id="smi08/ProArd")
+upload_file(path_or_fileobj="train_ofa_net_WPS.py",path_in_repo="",repo_id="smi08/ProArd")
+upload_file(path_or_fileobj="train_ofa_net.py",path_in_repo="",repo_id="smi08/ProArd")
+upload_file(path_or_fileobj="train_teacher_net.py",path_in_repo="",repo_id="smi08/ProArd")
+upload_file(path_or_fileobj="README.md",path_in_repo="",repo_id="smi08/ProArd")

sample_eval.py

@@ -0,0 +1,89 @@
+# Once for All: Train One Network and Specialize it for Efficient Deployment
+# Han Cai, Chuang Gan, Tianzhe Wang, Zhekai Zhang, Song Han
+# International Conference on Learning Representations (ICLR), 2020.
+
+import os
+import torch
+import argparse
+import sys    
+from proard.classification.data_providers.imagenet import ImagenetDataProvider
+from proard.classification.data_providers.cifar10 import Cifar10DataProvider
+from proard.classification.data_providers.cifar100 import Cifar100DataProvider
+from proard.classification.run_manager import ClassificationRunConfig, RunManager,DistributedRunManager
+from proard.model_zoo import DYN_net
+from proard.nas.accuracy_predictor import AccuracyDataset,AccuracyPredictor,ResNetArchEncoder,RobustnessPredictor,MobileNetArchEncoder,AccuracyRobustnessDataset,Accuracy_Robustness_Predictor
+
+parser = argparse.ArgumentParser()
+parser.add_argument(
+    "-p", "--path", help="The path of imagenet", type=str, default="/dataset/imagenet"
+)
+parser.add_argument("-g", "--gpu", help="The gpu(s) to use", type=str, default="all")
+parser.add_argument(
+    "-b",
+    "--batch-size",
+    help="The batch on every device for validation",
+    type=int,
+    default=128,
+)
+parser.add_argument("-j", "--workers", help="Number of workers", type=int, default=20)
+parser.add_argument(
+    "-n",
+    "--net",
+    metavar="DYNNET",
+    default="MBV3",
+    choices=[
+        "ResNet50",
+        "MBV3",
+        "ProxylessNASNet",
+        "MBV2"
+    ],
+    help="dynamic networks",
+)
+parser.add_argument(
+    "--dataset", type=str, default="cifar10" ,choices=["cifar10", "cifar100", "imagenet"]
+)
+parser.add_argument("--train_criterion", type=str, default="trades",choices=["trades","sat","mart","hat"])
+parser.add_argument(
+    "--robust_mode", type=bool, default=True
+)
+parser.add_argument(
+    "--WPS", type=bool, default=False
+)
+args = parser.parse_args()
+if args.gpu == "all":
+    device_list = range(torch.cuda.device_count())
+    args.gpu = ",".join(str(_) for _ in device_list)
+else:
+    device_list = [int(_) for _ in args.gpu.split(",")]
+os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
+args.batch_size = args.batch_size * max(len(device_list), 1)
+ImagenetDataProvider.DEFAULT_PATH = args.path
+
+run_config = ClassificationRunConfig(dataset= args.dataset, test_batch_size=args.batch_size, n_worker=args.workers,robust_mode=args.robust_mode)
+dyn_network = DYN_net(args.net,args.robust_mode,args.dataset, args.train_criterion ,pretrained=True,run_config=run_config,WPS=args.WPS)
+""" Randomly sample a sub-network, 
+    you can also manually set the sub-network using: 
+        dyn_network.set_active_subnet(ks=7, e=6, d=4) 
+"""
+# dyn_network.set_active_subnet(ks=3, e=3, d=2)
+# dyn_network.set_active_subnet(d=4,e=0.25,w=1) 
+import random 
+import numpy as np 
+random.seed(0)
+np.random.seed(0)
+acc1,rob1,acc2,rob2 =[],[],[],[]
+if args.net == "ResNet50":   
+   arch = ResNetArchEncoder(image_size_list=[224 if args.dataset == 'imagenet' else 32],depth_list=[0,1,2],expand_list=[0.2,0.25,0.35],width_mult_list=[0.65,0.8,1.0])
+else:
+   arch =  MobileNetArchEncoder (image_size_list=[224 if args.dataset == 'imagenet' else 32],depth_list=[2,3,4],expand_list=[3,4,6],ks_list=[3,5,7])
+print(arch)
+acc_data = AccuracyRobustnessDataset("./acc_rob_data_{}_{}_{}".format(args.dataset,args.net,args.train_criterion))
+train_loader, valid_loader, base_acc ,base_rob = acc_data.build_acc_data_loader(arch)
+for inputs, targets_acc, targets_rob in train_loader:
+    for i in range(len(targets_acc)):
+        acc1.append(targets_acc[i].item() * 100)
+        rob1.append(targets_rob[i].item() * 100)
+    
+np.save("./results/acc_mbv3.npy",np.array(acc1))
+np.save("./results/rob_mbv3.npy",np.array(rob1))
+

search_best.py

@@ -0,0 +1,273 @@
+import os
+import torch
+import argparse
+import torch.nn as nn 
+import numpy as np
+from tqdm.auto import tqdm
+from torch.utils.data import DataLoader
+import torch.nn as nn 
+import torch
+import random
+from torch import nn
+from torch.optim import *
+from torch.optim.lr_scheduler import *
+from torch.utils.data import DataLoader
+from torchprofile import profile_macs
+from torchvision.datasets import *
+from torchvision.transforms import *
+from proard.model_zoo import DYN_net
+from proard.nas.accuracy_predictor import AccuracyPredictor,ResNetArchEncoder,RobustnessPredictor,MobileNetArchEncoder,Accuracy_Robustness_Predictor
+from proard.nas.efficiency_predictor import ResNet50FLOPsModel,Mbv3FLOPsModel,ProxylessNASFLOPsModel
+from proard.nas.search_algorithm import EvolutionFinder,DynIndividual_mbv,DynIndividual_res,DynRandomSampler,DynProblem_mbv,DynProblem_res,DynSampling,individual_to_arch_res,individual_to_arch_mbv
+from utils.profile import trainable_param_num
+from pymoo.core.individual import Individual
+from pymoo.core.mutation import Mutation
+from pymoo.core.population import Population
+from pymoo.core.problem import Problem
+from pymoo.core.sampling import Sampling
+from pymoo.core.variable import Choice
+from pymoo.operators.crossover.ux import UniformCrossover
+from pymoo.operators.mutation.pm import PolynomialMutation
+from pymoo.operators.mutation.rm import ChoiceRandomMutation
+from pymoo.operators.selection.rnd import RandomSelection
+from pymoo.operators.selection.tournament import TournamentSelection
+from pymoo.algorithms.moo.nsga2 import NSGA2
+from pymoo.algorithms.moo.sms import SMSEMOA
+from pymoo.algorithms.moo.spea2 import SPEA2
+from pymoo.optimize import minimize
+from pymoo.termination import get_termination
+from pymoo.termination.default import DefaultMultiObjectiveTermination
+from pymoo.core.callback import Callback
+from pymoo.util.display.column import Column
+from pymoo.util.display.output import Output
+from proard.classification.run_manager import ClassificationRunConfig, RunManager
+parser = argparse.ArgumentParser()
+parser.add_argument(
+    "-p", "--path", help="The path of cifar10", type=str, default="/dataset/cifar10"
+)
+parser.add_argument("-g", "--gpu", help="The gpu(s) to use", type=str, default="all")
+parser.add_argument(
+    "-b",
+    "--batch-size",
+    help="The batch on every device for validation",
+    type=int,
+    default=100,
+)
+parser.add_argument("-j", "--workers", help="Number of workers", type=int, default=20)
+parser.add_argument(
+    "-n",
+    "--net",
+    metavar="DYNNET",
+    default="ResNet50",
+    choices=[
+        "ResNet50",
+        "MBV3",
+        "ProxylessNASNet",
+    ],
+    help="dynamic networks",
+)
+parser.add_argument(
+    "--dataset", type=str, default="cifar10" ,choices=["cifar10", "cifar100", "imagenet"]
+)
+parser.add_argument(
+    "--attack", type=str, default="linf-pgd" ,choices=['fgsm', 'linf-pgd', 'fgm', 'l2-pgd', 'linf-df', 'l2-df', 'linf-apgd', 'l2-apgd','squar_attack','autoattack','apgd_ce']
+)
+parser.add_argument("--train_criterion", type=str, default="trades",choices=["trades","sat","mart","hat"])
+parser.add_argument(
+    "--robust_mode", type=bool, default=True
+)
+args = parser.parse_args()
+if args.gpu == "all":
+    device_list = range(torch.cuda.device_count())
+    args.gpu = ",".join(str(_) for _ in device_list)
+else:
+    device_list = [int(_) for _ in args.gpu.split(",")]
+os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
+args.batch_size = args.batch_size * max(len(device_list), 1)
+run_config = ClassificationRunConfig(attack_type=args.attack, dataset= args.dataset, test_batch_size=args.batch_size, n_worker=args.workers,robust_mode=args.robust_mode)
+dyn_network = DYN_net(args.net,args.robust_mode,args.dataset,args.train_criterion, pretrained=True,run_config=run_config)
+if args.net == "ResNet50":
+    efficiency_predictor = ResNet50FLOPsModel(dyn_network)
+    arch = ResNetArchEncoder(image_size_list=[32],depth_list=[0,1,2],expand_list=[0.2,0.25,0.35],width_mult_list=[0.65,0.8,1.0])
+    accuracy_robustness_predictor = Accuracy_Robustness_Predictor(arch)
+    accuracy_robustness_predictor.load_state_dict(torch.load("./acc_rob_data_{}_{}_{}/src/model_acc_rob.pth".format(args.dataset,args.net,args.train_criterion)))
+elif args.net == "MBV3":
+    efficiency_predictor = Mbv3FLOPsModel(dyn_network)
+    arch = MobileNetArchEncoder(image_size_list=[32],depth_list=[2,3,4],expand_list=[3,4,6],ks_list=[3,5,7])
+    accuracy_robustness_predictor = Accuracy_Robustness_Predictor(arch)
+    accuracy_robustness_predictor.load_state_dict(torch.load("./acc_rob_data_{}_{}_{}/src/model_acc_rob.pth".format(args.dataset,args.net,args.train_criterion)))
+elif args.net == "ProxylessNASNet":
+    efficiency_predictor = ProxylessNASFLOPsModel(dyn_network)
+    arch = MobileNetArchEncoder(image_size_list=[32],depth_list=[2,3,4],expand_list=[3,4,6],width_mult_list=[3,5,7])
+    accuracy_robustness_predictor = Accuracy_Robustness_Predictor(arch)
+    accuracy_robustness_predictor.load_state_dict(torch.load("./acc_rob_data_{}_{}_{}/src/model_acc_rob.pth".format(args.dataset,args.net,args.train_criterion)))
+##### Test #################################################
+dyn_sampler = DynRandomSampler(arch, efficiency_predictor)
+# arch1, eff1 = dyn_sampler.random_sample()
+# arch2, eff2 = dyn_sampler.random_sample()
+# print(accuracy_predictor.predict_acc([arch1, arch2]))
+# print(arch1,eff1)
+##################################################
+
+""" Hyperparameters
+- P: size of the population in each generation (number of individuals)
+- N: number of generations to run the algorithm
+- mutate_prob: probability of gene mutation in the evolutionary search
+"""
+P = 100
+N = 100
+mutation_prob = 0.5
+
+
+# variables options
+if args.net == 'ResNet50':
+    search_space = {
+        'e': [0.2, 0.25, 0.35],
+        'd': [0, 1, 2],
+        'w': [0 ,1 ,2],
+        'image_size': [32]
+    }
+else:
+    search_space = {
+        'ks': [3, 5, 7],
+        'e': [3, 4, 6],
+        'd': [2, 3, 4],
+        'image_size': [32]
+    }
+
+#----------------------------
+# units
+num_blocks = arch.max_n_blocks
+num_stages = arch.n_stage
+Flops_constraints = 1600
+if args.net == "ResNet50":
+    problem = DynProblem_res(efficiency_predictor, accuracy_robustness_predictor, num_blocks, num_stages, search_space,Flops_constraints)
+else:
+    problem = DynProblem_mbv(efficiency_predictor, accuracy_robustness_predictor, num_blocks, num_stages, search_space,Flops_constraints)
+
+
+
+
+
+mutation_rc = ChoiceRandomMutation(prob=1.0, prob_var=0.1)
+crossover_ux = UniformCrossover(prob=1.0)
+# selection_tournament = TournamentSelection(
+#     func_comp=accuracy_predictor.predict_acc, 
+#     pressure=2
+# )
+termination_default = DefaultMultiObjectiveTermination(
+    xtol=1e-8, cvtol=1e-6, ftol=0.0025, period=30, n_max_gen=1000, n_max_evals=100000
+)
+termination_gen = get_termination("n_gen", N)
+np.random.seed(42)
+random.seed(42)
+if args.net=="ResNet50":
+    init_pop = Population(individuals=[DynIndividual_res(dyn_sampler.random_sample(), accuracy_robustness_predictor) for _ in range(P)])
+else:
+    init_pop = Population(individuals=[DynIndividual_mbv(dyn_sampler.random_sample(), accuracy_robustness_predictor) for _ in range(P)])
+
+algorithm = NSGA2(
+        pop_size=P,
+        sampling=DynSampling(),
+        # selection=selection_tournament,
+        crossover=crossover_ux,
+        mutation=mutation_rc,
+        # mutation=mutation_pm,
+        # survival=RankAndCrowdingSurvival(),
+        # output=MultiObjectiveOutput(),
+        #    **kwargs
+    )
+res_nsga2 = minimize(
+    problem,
+    algorithm,
+    termination=termination_gen,
+    seed=1,
+    #verbose=True,
+    verbose=False,
+    save_history=True,
+)
+# print(100-res_nsga2.history[99].pop.get('F')[:,0],100-res_nsga2.history[99].pop.get('F')[:,1])
+# a = individual_to_arch_res(res_nsga2.pop.get('X'),num_blocks)[0]
+# # print(a)
+# # a['d'][3]  = int(a['d'][3])
+# a['d'][4]  = int(a['d'][4])
+# dyn_network.set_active_subnet(**a)
+# subnet = dyn_network.get_active_subnet(preserve_weight=True)
+# run_manager = RunManager(".tmp/eval_subnet", subnet, run_config, init=False)
+# run_config.data_provider.assign_active_img_size(32)
+# run_manager.reset_running_statistics(net=subnet)
+
+# print("Test random subnet:")
+# # print(subnet.module_str)
+
+# loss, (top1, top5,robust1,robust5) = run_manager.validate(net=subnet,is_test=True)
+# print("Results: loss=%.5f,\t top1=%.1f,\t top5=%.1f,\t robust1=%.1f,\t robust5=%.1f" % (loss, top1, top5,robust1,robust5))
+
+
+np.savetxt("./results/acc_gen0.csv", 100-res_nsga2.history[0].pop.get('F')[:,0], delimiter=",") 
+
+np.savetxt("./results/acc_gen99.csv", 100-res_nsga2.history[99].pop.get('F')[:,0], delimiter=",")   
+np.savetxt("./results/rob_gen0.csv", 100-res_nsga2.history[0].pop.get('F')[:,1], delimiter=",") 
+
+np.savetxt("./results/rob_gen99.csv", 100-res_nsga2.history[99].pop.get('F')[:,1], delimiter=",")
+np.savetxt("./results/flops_gen99.csv", res_nsga2.history[99].pop.get('G'), delimiter=",")
+
+# np.savetxt("./results/robs.csv", np.array(robs), delimiter=",") 
+
+from matplotlib import pyplot as plt
+from matplotlib.ticker import FormatStrFormatter
+from matplotlib.ticker import AutoMinorLocator, MultipleLocator
+# NSGA-II population progression
+x_min, x_max, y_min, y_max = 80, 93, 47, 56
+ax_limits = [x_min, x_max, y_min, y_max]
+#-------------------------------------------------
+# plot
+fig, ax = plt.subplots(dpi=600)
+gen0 = 0
+gen1 = 99
+print(100-res_nsga2.history[gen1].pop.get('F')[:,0], 100 - res_nsga2.history[gen1].pop.get('F')[:,1])
+# gen2 = 99
+# print(res_nsga2.history[gen0].pop.get('F')[:,0],res_nsga2.history[gen0].pop.get('F')[:,1]  )
+ax.plot(100-res_nsga2.history[gen0].pop.get('F')[:,0], 100 - res_nsga2.history[gen0].pop.get('F')[:,1]  , 'o', label=f'Population at generation #{gen0+1}', color='red',    alpha=0.5)
+ax.plot(100-res_nsga2.history[gen1].pop.get('F')[:,0], 100 - res_nsga2.history[gen1].pop.get('F')[:,1] , 'o', label=f'Population at generation #{gen1+1}', color='green',  alpha=0.5)
+# ax.plot(res_nsga2.history[gen2].pop.get('F')[:,0], 100 - res_nsga2.history[gen2].pop.get('F')[:,1], 'o', label=f'Population at generation #{gen2+1}', color='orange', alpha=0.5)
+# ax.plot(res_nsga2.history[gen3].pop.get('F')[:,0], 100 - res_nsga2.history[gen3].pop.get('F')[:,1], 'o', label=f'Population at generation #{gen3+1}', color='blue',   alpha=0.5)
+#-------------------------------------------------
+# text
+ax.grid(True, linestyle=':')
+ax.set_xlabel('Accuracy (%)')
+ax.set_ylabel('Robustness (%)')
+ax.set_title('NSGA-II solutions progression For Fixed number of FLOPs'),
+ax.legend()
+#-------------------------------------------------
+# x-axis
+ax.xaxis.set_major_locator(MultipleLocator(1))
+ax.xaxis.set_minor_locator(MultipleLocator(1))
+# y-axis
+ax.yaxis.set_major_locator(MultipleLocator(1))
+ax.yaxis.set_minor_locator(MultipleLocator(1))
+# ax.yaxis.set_major_formatter(FormatStrFormatter('%.1f'))
+ax.set(xlim=(ax_limits[0], ax_limits[1]), ylim=(ax_limits[2], ax_limits[3]))
+#-------------------------------------------------
+plt.savefig('nsga2_pop_progression_debug.png')
+fig.set_dpi(100)
+# plt.close(fig)
+
+
+
+# plt.show()
+
+
+# finder  = EvolutionFinder(efficiency_predictor,accuracy_predictor,Robustness_predictor)
+# valid_constraint_range = 800
+# best_valids, best_info = finder.run_evolution_search(constraint=valid_constraint_range,verbose=True)
+# print(efficiency_predictor.get_efficiency(best_info[2]))
+# dyn_network.set_active_subnet(best_info[2]['d'],best_info[2]['e'],best_info[2]['w'])
+# subnet = dyn_network.get_active_subnet(preserve_weight=True)
+# run_config = CifarRunConfig_robust(test_batch_size=args.batch_size, n_worker=args.workers)
+# run_manager = RunManager_robust(".tmp/eval_subnet", subnet, run_config, init=False)
+# run_config.data_provider.assign_active_img_size(32)
+# run_manager.reset_running_statistics(net=subnet)
+# loss, (top1, top5,robust1,robust5) = run_manager.validate(net=subnet)
+# print("Results: loss=%.5f,\t top1=%.1f,\t top5=%.1f,\t robust1=%.1f,\t robust5=%.1f" % (loss, top1, top5,robust1,robust5))
+# print("number of parameter={}M".format(trainable_param_num(subnet)))

train_ofa_net.py

@@ -0,0 +1,558 @@
+# Once for All: Train One Network and Specialize it for Efficient Deployment
+# Han Cai, Chuang Gan, Tianzhe Wang, Zhekai Zhang, Song Han
+# International Conference on Learning Representations (ICLR), 2020.
+
+import argparse
+import numpy as np
+import os
+import random
+
+# using for distributed training
+import horovod.torch as hvd
+import torch
+
+
+from proard.classification.elastic_nn.modules.dynamic_op import (
+    DynamicSeparableConv2d,
+)
+from proard.classification.elastic_nn.networks import DYNMobileNetV3,DYNProxylessNASNets,DYNResNets,DYNProxylessNASNets_Cifar,DYNMobileNetV3_Cifar,DYNResNets_Cifar
+from proard.classification.run_manager import DistributedClassificationRunConfig
+from proard.classification.run_manager.distributed_run_manager import (
+    DistributedRunManager
+)
+from proard.utils import download_url, MyRandomResizedCrop
+from proard.classification.elastic_nn.training.progressive_shrinking import load_models
+
+parser = argparse.ArgumentParser()
+parser.add_argument(
+    "--task",
+    type=str,
+    default="expand",
+    choices=[
+        "kernel", # for architecture except ResNet
+        "depth",
+        "expand",
+        "width", # only for ResNet
+    ],
+)
+parser.add_argument("--phase", type=int, default=2, choices=[1, 2])
+parser.add_argument("--resume", action="store_true")
+parser.add_argument("--model_name", type=str, default="MBV2", choices=["ResNet50", "MBV3", "ProxylessNASNet","MBV2"])
+parser.add_argument("--dataset", type=str, default="cifar100", choices=["cifar10", "cifar100", "imagenet"])
+parser.add_argument("--robust_mode", type=bool, default=True)
+parser.add_argument("--epsilon", type=float, default=0.031)
+parser.add_argument("--num_steps", type=int, default=10)
+parser.add_argument("--step_size", type=float, default=0.0078) 
+parser.add_argument("--clip_min", type=int, default=0)
+parser.add_argument("--clip_max", type=int, default=1)
+parser.add_argument("--const_init", type=bool, default=False)
+parser.add_argument("--beta", type=float, default=6.0)
+parser.add_argument("--distance", type=str, default="l_inf",choices=["l_inf","l2"])
+parser.add_argument("--train_criterion", type=str, default="trades",choices=["trades","sat","mart","hat"])
+parser.add_argument("--test_criterion", type=str, default="ce",choices=["ce"])
+parser.add_argument("--kd_criterion", type=str, default="rslad",choices=["ard","rslad","adaad"])
+parser.add_argument("--attack_type", type=str, default="linf-pgd",choices=['fgsm', 'linf-pgd', 'fgm', 'l2-pgd', 'linf-df', 'l2-df', 'linf-apgd', 'l2-apgd','squar_attack','autoattack','apgd_ce'])
+
+args = parser.parse_args()
+if args.model_name == "ResNet50":
+    args.ks_list = "3"
+    if args.task == "width":
+        if args.robust_mode:
+            args.path = "exp/robust/"+ args.dataset + '/' + args.model_name +'/' + args.train_criterion +"/normal2width"
+        else:
+            args.path = "exp/"+ args.dataset + '/' +args.model_name +'/' + args.train_criterion +"/normal2width"    
+        args.dynamic_batch_size = 1
+        args.n_epochs = 120
+        args.base_lr = 3e-2
+        args.warmup_epochs = 5
+        args.warmup_lr = -1
+        args.width_mult_list = "0.65,0.8,1.0"
+        args.expand_list = "0.35"
+        args.depth_list = "2"
+    elif args.task == "depth":
+        if args.robust_mode:
+            args.path = "exp/robust/"+ args.dataset + '/'  + args.model_name +'/' + args.train_criterion +"/width2width_depth/phase%d" % args.phase
+        else:
+            args.path = "exp/"+ args.dataset + '/'  + args.model_name +'/' + args.train_criterion +"/width2width_depth/phase%d" % args.phase     
+        args.dynamic_batch_size = 2
+        if args.phase == 1:
+            args.n_epochs = 25
+            args.base_lr = 2.5e-3
+            args.warmup_epochs = 0
+            args.warmup_lr = -1
+            args.width_mult_list = "0.65,0.8,1.0"
+            args.expand_list ="0.35"
+            args.depth_list = "1,2"
+        else:
+            args.n_epochs = 120
+            args.base_lr = 7.5e-3
+            args.warmup_epochs = 5
+            args.warmup_lr = -1
+            args.width_mult_list = "0.65,0.8,1.0"
+            args.expand_list = "0.35"
+            args.depth_list = "0,1,2"
+    elif args.task == "expand":
+        if args.robust_mode : 
+            args.path = "exp/robust/"+ args.dataset + '/' + args.model_name +'/' + args.train_criterion +"/width_depth2width_depth_width/phase%d" % args.phase 
+        else:
+            args.path = "exp/"+ args.dataset + '/' + args.model_name +'/' + args.train_criterion +"/width_depth2width_depth_width/phase%d" % args.phase    
+        args.dynamic_batch_size = 4
+        if args.phase == 1:
+            args.n_epochs = 25
+            args.base_lr = 2.5e-3
+            args.warmup_epochs = 0
+            args.warmup_lr = -1
+            args.width_mult_list = "0.65,0.8,1.0"
+            args.expand_list = "0.25,0.35"
+            args.depth_list = "0,1,2"
+        else:
+            args.n_epochs = 120
+            args.base_lr = 7.5e-3
+            args.warmup_epochs = 5
+            args.warmup_lr = -1
+            args.width_mult_list = "0.65,0.8,1.0"
+            args.expand_list = "0.2,0.25,0.35"
+            args.depth_list = "0,1,2"
+    else:
+        raise NotImplementedError
+else:
+    args.width_mult_list = "1.0"    
+    if args.task == "kernel":
+        if args.robust_mode:
+            args.path = "exp/robust/"+ args.dataset + '/' +  args.model_name +'/' + args.train_criterion +"/normal2kernel"
+        else:    
+            args.path = "exp/"+ args.dataset + '/' +  args.model_name +'/' + args.train_criterion +"/normal2kernel"
+        args.dynamic_batch_size = 1
+        args.n_epochs = 120
+        args.base_lr = 3e-2
+        args.warmup_epochs = 5
+        args.warmup_lr = -1
+        args.ks_list = "3,5,7"
+        args.expand_list = "6"
+        args.depth_list = "4"
+    elif args.task == "depth":
+        if args.robust_mode : 
+            args.path = "exp/robust/"+args.dataset + '/' + args.model_name +'/' + args.train_criterion +"/kernel2kernel_depth/phase%d" % args.phase
+        else:
+            args.path = "exp/"+args.dataset + '/' + args.model_name +'/' + args.train_criterion +"/kernel2kernel_depth/phase%d" % args.phase    
+        args.dynamic_batch_size = 2
+        if args.phase == 1:
+            args.n_epochs = 25
+            args.base_lr = 2.5e-3
+            args.warmup_epochs = 0
+            args.warmup_lr = -1
+            args.ks_list = "3,5,7"
+            args.expand_list = "6"
+            args.depth_list = "3,4"
+        else:
+            args.n_epochs = 120
+            args.base_lr = 7.5e-3
+            args.warmup_epochs = 5
+            args.warmup_lr = -1
+            args.ks_list = "3,5,7"
+            args.expand_list = "6"
+            args.depth_list = "2,3,4"
+    elif args.task == "expand":
+        if args.robust_mode:
+            args.path = "exp/robust/"+ args.dataset + '/' + args.model_name +'/' + args.train_criterion +"/kernel_depth2kernel_depth_width/phase%d" % args.phase
+        else:
+            args.path = "exp/"+ args.dataset + '/' + args.model_name +  '/' + args.train_criterion + "/kernel_depth2kernel_depth_width/phase%d" % args.phase    
+        args.dynamic_batch_size = 4
+        if args.phase == 1:
+            args.n_epochs = 25
+            args.base_lr = 2.5e-3
+            args.warmup_epochs = 0
+            args.warmup_lr = -1
+            args.ks_list = "3,5,7"
+            args.expand_list = "4,6"
+            args.depth_list = "2,3,4"
+        else:
+            args.n_epochs = 120
+            args.base_lr = 7.5e-3
+            args.warmup_epochs = 5
+            args.warmup_lr = -1
+            args.ks_list = "3,5,7"
+            args.expand_list = "3,4,6"
+            args.depth_list = "2,3,4"
+    else:
+        raise NotImplementedError
+args.manual_seed = 0
+
+args.lr_schedule_type = "cosine"
+
+args.base_batch_size = 64
+args.valid_size = 64
+
+args.opt_type = "sgd"
+args.momentum = 0.9
+args.no_nesterov = False
+args.weight_decay = 3e-5
+args.label_smoothing = 0.1
+args.no_decay_keys = "bn#bias"
+args.fp16_allreduce = False
+
+args.model_init = "he_fout"
+args.validation_frequency = 1
+args.print_frequency = 10
+
+args.n_worker = 8
+args.resize_scale = 0.08
+args.distort_color = "tf"
+if args.dataset == "imagenet":  
+    args.image_size = "128,160,192,224"
+else:
+    args.image_size = "32"    
+args.continuous_size = True
+args.not_sync_distributed_image_size = False
+
+args.bn_momentum = 0.1
+args.bn_eps = 1e-5
+args.dropout = 0.1
+args.base_stage_width = "google"
+
+
+args.dy_conv_scaling_mode = 1
+args.independent_distributed_sampling = False
+
+args.kd_ratio = 1.0
+args.kd_type = "ce"
+
+
+if __name__ == "__main__":
+    os.makedirs(args.path, exist_ok=True)
+
+    # Initialize Horovod
+    hvd.init()
+    # Pin GPU to be used to process local rank (one GPU per process)
+    torch.cuda.set_device(hvd.local_rank())
+    if args.robust_mode:
+        args.teacher_path = 'exp/robust/teacher/' + args.dataset + '/' +  args.model_name + '/' + args.train_criterion + "/checkpoint/model_best.pth.tar"
+    else:
+        args.teacher_path = 'exp/teacher/' + args.dataset + '/' +  args.model_name +'/' + args.train_criterion + "/checkpoint/model_best.pth.tar"
+    num_gpus = hvd.size()
+
+    torch.manual_seed(args.manual_seed)
+    torch.cuda.manual_seed_all(args.manual_seed)
+    np.random.seed(args.manual_seed)
+    random.seed(args.manual_seed)
+
+    # image size
+    args.image_size = [int(img_size) for img_size in args.image_size.split(",")]
+    if len(args.image_size) == 1:
+        args.image_size = args.image_size[0]
+    MyRandomResizedCrop.CONTINUOUS = args.continuous_size
+    MyRandomResizedCrop.SYNC_DISTRIBUTED = not args.not_sync_distributed_image_size
+
+    # build run config from args
+    args.lr_schedule_param = None
+    args.opt_param = {
+        "momentum": args.momentum,
+        "nesterov": not args.no_nesterov,
+    }
+    args.init_lr = args.base_lr * num_gpus  # linearly rescale the learning rate
+    if args.warmup_lr < 0:
+        args.warmup_lr = args.base_lr
+    args.train_batch_size = args.base_batch_size
+    args.test_batch_size = args.base_batch_size * 4
+    run_config = DistributedClassificationRunConfig(
+        **args.__dict__, num_replicas=num_gpus, rank=hvd.rank()
+    )
+
+    # print run config information
+    if hvd.rank() == 0:
+        print("Run config:")
+        for k, v in run_config.config.items():
+            print("\t%s: %s" % (k, v))
+
+    if args.dy_conv_scaling_mode == -1:
+        args.dy_conv_scaling_mode = None
+    DynamicSeparableConv2d.KERNEL_TRANSFORM_MODE = args.dy_conv_scaling_mode
+
+    # build net from args
+    args.width_mult_list = [
+        float(width_mult) for width_mult in args.width_mult_list.split(",")
+    ]
+    args.ks_list = [int(ks) for ks in args.ks_list.split(",")]
+    if args.model_name == "ResNet50":
+        args.expand_list = [float(e) for e in args.expand_list.split(",")]
+    else:
+        args.expand_list = [int(e) for e in args.expand_list.split(",")]    
+    args.depth_list = [int(d) for d in args.depth_list.split(",")]
+
+    args.width_mult_list = (
+        args.width_mult_list[0]
+        if len(args.width_mult_list) == 1
+        else args.width_mult_list
+    )
+
+    if args.model_name == "ResNet50":
+        if args.dataset == "cifar10" or args.dataset == "cifar100":
+            net = DYNResNets_Cifar( n_classes=run_config.data_provider.n_classes,
+                bn_param=(args.bn_momentum, args.bn_eps),
+                dropout_rate=args.dropout,
+                depth_list=args.depth_list,
+                expand_ratio_list=args.expand_list,
+                width_mult_list=args.width_mult_list,)
+        else:
+            net = DYNResNets( n_classes=run_config.data_provider.n_classes,
+                bn_param=(args.bn_momentum, args.bn_eps),
+                dropout_rate=args.dropout,
+                depth_list=args.depth_list,
+                expand_ratio_list=args.expand_list,
+                width_mult_list=args.width_mult_list,)   
+    elif args.model_name == "MBV3":  
+        if args.dataset == "cifar10" or args.dataset == "cifar100":  
+            net = DYNMobileNetV3_Cifar(n_classes=run_config.data_provider.n_classes,bn_param=(args.bn_momentum,args.bn_eps),
+                                dropout_rate= args.dropout, ks_list=args.ks_list , expand_ratio_list= args.expand_list , depth_list= args.depth_list,width_mult=args.width_mult_list)
+        else:
+            net = DYNMobileNetV3(n_classes=run_config.data_provider.n_classes,bn_param=(args.bn_momentum,args.bn_eps),
+                                dropout_rate= args.dropout, ks_list=args.ks_list , expand_ratio_list= args.expand_list , depth_list= args.depth_list,width_mult=args.width_mult_list)    
+    elif args.model_name == "ProxylessNASNet": 
+        if args.dataset == "cifar10" or args.dataset == "cifar100":     
+            net = DYNProxylessNASNets_Cifar(n_classes=run_config.data_provider.n_classes,bn_param=(args.bn_momentum,args.bn_eps),
+                                dropout_rate= args.dropout, ks_list=args.ks_list , expand_ratio_list= args.expand_list , depth_list= args.depth_list,width_mult=args.width_mult_list)
+        else:
+            net = DYNProxylessNASNets(n_classes=run_config.data_provider.n_classes,bn_param=(args.bn_momentum,args.bn_eps),
+                                dropout_rate= args.dropout, ks_list=args.ks_list , expand_ratio_list= args.expand_list , depth_list= args.depth_list,width_mult=args.width_mult_list)   
+    elif args.model_name == "MBV2": 
+        if args.dataset == "cifar10" or args.dataset == "cifar100":     
+            net = DYNProxylessNASNets_Cifar(n_classes=run_config.data_provider.n_classes,bn_param=(args.bn_momentum,args.bn_eps),
+                                dropout_rate= args.dropout, ks_list=args.ks_list , expand_ratio_list= args.expand_list , depth_list= args.depth_list,width_mult=args.width_mult_list,base_stage_width=args.base_stage_width)
+        else:
+            net = DYNProxylessNASNets(n_classes=run_config.data_provider.n_classes,bn_param=(args.bn_momentum,args.bn_eps),
+                                dropout_rate= args.dropout, ks_list=args.ks_list , expand_ratio_list= args.expand_list , depth_list= args.depth_list,width_mult=args.width_mult_list,base_stage_width=args.base_stage_width)             
+    else: 
+        raise NotImplementedError  
+    # teacher model
+    if args.kd_ratio > 0:
+        
+        if args.model_name =="ResNet50":
+            if args.dataset == "cifar10" or args.dataset == "cifar100":
+                args.teacher_model = DYNResNets_Cifar(
+                    n_classes=run_config.data_provider.n_classes,
+                    bn_param=(args.bn_momentum, args.bn_eps),
+                    dropout_rate=args.dropout,
+                    depth_list=[2],
+                    expand_ratio_list=[0.35],
+                    width_mult_list=[1.0],
+                )
+            else:
+                args.teacher_model = DYNResNets(
+                    n_classes=run_config.data_provider.n_classes,
+                    bn_param=(args.bn_momentum, args.bn_eps),
+                    dropout_rate=args.dropout,
+                    depth_list=[2],
+                    expand_ratio_list=[0.35],
+                    width_mult_list=[1.0],
+                )   
+        elif args.model_name =="MBV3":    
+            if args.dataset == "cifar10" or args.dataset == "cifar100":
+                args.teacher_model = DYNMobileNetV3_Cifar(
+                    n_classes=run_config.data_provider.n_classes,
+                    bn_param=(args.bn_momentum, args.bn_eps),
+                    dropout_rate=0,
+                    width_mult=1.0,
+                    ks_list=[7],
+                    expand_ratio_list=[6],
+                    depth_list=[4]
+                )
+            else:
+                args.teacher_model = DYNMobileNetV3(
+                    n_classes=run_config.data_provider.n_classes,
+                    bn_param=(args.bn_momentum, args.bn_eps),
+                    dropout_rate=0,
+                    width_mult=1.0,
+                    ks_list=[7],
+                    expand_ratio_list=[6],
+                    depth_list=[4]
+                )    
+        elif args.model_name == "ProxylessNASNet":
+            if args.dataset == "cifar10" or args.dataset == "cifar100":
+                args.teacher_model  = DYNProxylessNASNets_Cifar(n_classes=run_config.data_provider.n_classes,
+                    bn_param=(args.bn_momentum, args.bn_eps),
+                    dropout_rate=0,
+                    width_mult=1.0,
+                    ks_list=[7],
+                    expand_ratio_list=[6],
+                    depth_list=[4])   
+            else:
+                args.teacher_model  = DYNProxylessNASNets(n_classes=run_config.data_provider.n_classes,
+                    bn_param=(args.bn_momentum, args.bn_eps),
+                    dropout_rate=0,
+                    width_mult=1.0,
+                    ks_list=[7],
+                    expand_ratio_list=[6],
+                    depth_list=[4]) 
+        elif args.model_name == "MBV2":
+            if args.dataset == "cifar10" or args.dataset == "cifar100":
+                args.teacher_model  = DYNProxylessNASNets_Cifar(n_classes=run_config.data_provider.n_classes,
+                    bn_param=(args.bn_momentum, args.bn_eps),
+                    dropout_rate=0,
+                    width_mult=1.0,
+                    ks_list=[7],
+                    expand_ratio_list=[6],
+                    depth_list=[4],base_stage_width=args.base_stage_width)   
+            else:
+                args.teacher_model  = DYNProxylessNASNets(n_classes=run_config.data_provider.n_classes,
+                    bn_param=(args.bn_momentum, args.bn_eps),
+                    dropout_rate=0,
+                    width_mult=1.0,
+                    ks_list=[7],
+                    expand_ratio_list=[6],
+                    depth_list=[4],base_stage_width=args.base_stage_width)                 
+        args.teacher_model.cuda()
+    
+    """ Distributed RunManager """
+    # Horovod: (optional) compression algorithm.
+    compression = hvd.Compression.fp16 if args.fp16_allreduce else hvd.Compression.none
+    distributed_run_manager = DistributedRunManager(
+        args.path,
+        net,
+        run_config,
+        compression,
+        backward_steps=args.dynamic_batch_size,
+        is_root=(hvd.rank() == 0),
+    )
+    distributed_run_manager.save_config()
+    # hvd broadcast
+    distributed_run_manager.broadcast()
+
+    # load teacher net weights
+    if args.kd_ratio > 0:
+        load_models(
+            distributed_run_manager, args.teacher_model, model_path=args.teacher_path
+        )
+
+    # training
+    from proard.classification.elastic_nn.training.progressive_shrinking import (
+        validate,
+        train,
+    )
+    if args.model_name =="ResNet50":
+        validate_func_dict = {
+            "image_size_list": {224 if args.dataset == "imagenet" else 32}
+            if isinstance(args.image_size, int)
+            else sorted({160, 224}),
+            "width_mult_list": sorted({min(args.width_mult_list), max(args.width_mult_list)}),
+            "expand_ratio_list": sorted({min(args.expand_list), max(args.expand_list)}),
+            "depth_list": sorted({min(net.depth_list), max(net.depth_list)}),
+        }
+    else:
+        validate_func_dict = {
+            "image_size_list": {224 if args.dataset == "imagenet" else 32}
+            if isinstance(args.image_size, int)
+            else sorted({160, 224}),
+            "width_mult_list": [1.0],
+            "ks_list": sorted({min(args.ks_list), max(args.ks_list)}),
+            "expand_ratio_list": sorted({min(args.expand_list), max(args.expand_list)}),
+            "depth_list": sorted({min(net.depth_list), max(net.depth_list)}),
+        }  
+
+    if args.task == "width":
+        from proard.classification.elastic_nn.training.progressive_shrinking import (
+            train_elastic_width_mult,
+        )
+        if distributed_run_manager.start_epoch == 0:
+            if args.robust_mode:
+                args.dyn_checkpoint_path ='exp/robust/teacher/' +args.dataset + '/' +  args.model_name +'/' + args.train_criterion + "/checkpoint/model_best.pth.tar"
+            else:
+                args.dyn_checkpoint_path ='exp/teacher/' +args.dataset + '/' +  args.model_name +'/' + args.train_criterion + "/checkpoint/model_best.pth.tar"   
+            load_models(
+                distributed_run_manager,
+                distributed_run_manager.net,
+                args.dyn_checkpoint_path,
+            )
+            distributed_run_manager.write_log(
+                "%.3f\t%.3f\t%.3f\t%.3f\t%.3f\t%s"
+                % validate(distributed_run_manager, is_test=True, **validate_func_dict),
+                "valid",
+            )
+        else:
+            assert args.resume
+        train_elastic_width_mult (train,distributed_run_manager,args,validate_func_dict)    
+
+
+
+    elif args.task == "kernel":
+        validate_func_dict["ks_list"] = sorted(args.ks_list)
+        if distributed_run_manager.start_epoch == 0:
+            if args.robust_mode:
+                args.dyn_checkpoint_path ='exp/robust/teacher/' + args.dataset + '/' +  args.model_name +'/' + args.train_criterion + "/checkpoint/model_best.pth.tar"
+            else: 
+                args.dyn_checkpoint_path ='exp/teacher/' + args.dataset + '/' +  args.model_name +'/' + args.train_criterion + "/checkpoint/model_best.pth.tar"    
+            load_models(
+                distributed_run_manager,
+                distributed_run_manager.net,
+                args.dyn_checkpoint_path,
+            )
+            distributed_run_manager.write_log(
+               "%.3f\t%.3f\t%.3f\t%.3f\t%.3f\t%s"
+                % validate(distributed_run_manager, is_test=True, **validate_func_dict),
+                "valid",
+            )
+        else:
+            assert args.resume
+        train(
+            distributed_run_manager,
+            args,
+            lambda _run_manager, epoch, is_test: validate(
+                _run_manager, epoch, is_test, **validate_func_dict
+            ),
+        )
+    elif args.task == "depth":
+        from proard.classification.elastic_nn.training.progressive_shrinking import (
+            train_elastic_depth,
+        )
+        if args.robust_mode:
+            if args.model_name =="ResNet50":
+                if args.phase == 1:
+                    args.dyn_checkpoint_path =  "exp/robust/"+ args.dataset + '/'+ args.model_name +'/' + args.train_criterion +"/normal2width" +"/checkpoint/model_best.pth.tar"
+                else:
+                    args.dyn_checkpoint_path = "exp/robust/"+ args.dataset + '/' + args.model_name +'/' + args.train_criterion +"/width2width_depth/phase1" + "/checkpoint/model_best.pth.tar"
+            else:
+                if args.phase == 1:
+                    args.dyn_checkpoint_path =  "exp/robust/"+ args.dataset + '/' +  args.model_name +'/' + args.train_criterion +"/normal2kernel" +"/checkpoint/model_best.pth.tar"
+                else:
+                    args.dyn_checkpoint_path = "exp/robust/"+ args.dataset + '/' + args.model_name +'/' + args.train_criterion +"/kernel2kernel_depth/phase1" + "/checkpoint/model_best.pth.tar"     
+        else :
+            if args.model_name =="ResNet50":
+                if args.phase == 1:
+                    args.dyn_checkpoint_path =  "exp/"+ args.dataset + '/'+ args.model_name +'/' + args.train_criterion +"/normal2width" +"/checkpoint/model_best.pth.tar"
+                else:
+                    args.dyn_checkpoint_path = "exp/"+ args.dataset + '/' + args.model_name +'/' + args.train_criterion +"/width2width_depth/phase1" + "/checkpoint/model_best.pth.tar"
+            else:
+                if args.phase == 1:
+                    args.dyn_checkpoint_path =  "exp/"+ args.dataset + '/' +  args.model_name +'/' + args.train_criterion +"/normal2kernel" +"/checkpoint/model_best.pth.tar"
+                else:
+                    args.dyn_checkpoint_path = "exp/"+ args.dataset + '/' + args.model_name +'/' + args.train_criterion +"/kernel2kernel_depth/phase1" + "/checkpoint/model_best.pth.tar"                  
+        train_elastic_depth(train, distributed_run_manager, args, validate_func_dict)
+    elif args.task == "expand":
+        from proard.classification.elastic_nn.training.progressive_shrinking import (
+            train_elastic_expand,
+        )
+        if args.robust_mode : 
+            if args.model_name =="ResNet50":
+                if args.phase == 1:
+                    args.dyn_checkpoint_path =  "exp/robust/"+ args.dataset + '/'+ args.model_name +'/' + args.train_criterion +"/width2width_depth/phase2" + "/checkpoint/model_best.pth.tar"
+                else:
+                    args.dyn_checkpoint_path = "exp/robust/"+ args.dataset + '/'+ args.model_name +'/' + args.train_criterion +"/width_depth2width_depth_width/phase1" + "/checkpoint/model_best.pth.tar"
+            else:
+                if args.phase == 1:
+                    args.dyn_checkpoint_path =  "exp/robust/"+ args.dataset + '/'+ args.model_name +'/' + args.train_criterion +"/kernel2kernel_depth/phase2" + "/checkpoint/model_best.pth.tar"  
+                else:
+                    args.dyn_checkpoint_path = "exp/robust/"+ args.dataset + '/'+ args.model_name +'/' + args.train_criterion +"/kernel_depth2kernel_depth_width/phase1" +  "/checkpoint/model_best.pth.tar" 
+        else:
+            if args.model_name =="ResNet50":
+                if args.phase == 1:
+                    args.dyn_checkpoint_path =  "exp/"+ args.dataset + '/'+ args.model_name +'/' + args.train_criterion +"/width2width_depth/phase2" + "/checkpoint/model_best.pth.tar"
+                else:
+                    args.dyn_checkpoint_path = "exp/"+ args.dataset + '/'+ args.model_name +'/' + args.train_criterion +"/width_depth2width_depth_width/phase1" + "/checkpoint/model_best.pth.tar"
+            else:
+                if args.phase == 1:
+                    args.dyn_checkpoint_path =  "exp/"+ args.dataset + '/'+ args.model_name +'/' + args.train_criterion +"/kernel2kernel_depth/phase2" + "/checkpoint/model_best.pth.tar"  
+                else:
+                    args.dyn_checkpoint_path = "exp/"+ args.dataset + '/'+ args.model_name +'/' + args.train_criterion +"/kernel_depth2kernel_depth_width/phase1" +  "/checkpoint/model_best.pth.tar" 
+
+        train_elastic_expand(train, distributed_run_manager, args, validate_func_dict)
+    else:
+        raise NotImplementedError

train_ofa_net_WPS.py

@@ -0,0 +1,572 @@
+# Once for All: Train One Network and Specialize it for Efficient Deployment
+# Han Cai, Chuang Gan, Tianzhe Wang, Zhekai Zhang, Song Han
+# International Conference on Learning Representations (ICLR), 2020.
+
+import argparse
+import numpy as np
+import os
+import random
+
+# using for distributed training
+import horovod.torch as hvd
+import torch
+
+
+from proard.classification.elastic_nn.modules.dynamic_op import (
+    DynamicSeparableConv2d,
+)
+from proard.classification.elastic_nn.networks import DYNMobileNetV3,DYNProxylessNASNets,DYNResNets,DYNProxylessNASNets_Cifar,DYNMobileNetV3_Cifar,DYNResNets_Cifar
+from proard.classification.run_manager import DistributedClassificationRunConfig
+from proard.classification.run_manager.distributed_run_manager import (
+    DistributedRunManager
+)
+from proard.utils import download_url, MyRandomResizedCrop
+from proard.classification.elastic_nn.training.progressive_shrinking import (
+    load_models,
+)
+
+parser = argparse.ArgumentParser()
+parser.add_argument(
+    "--task",
+    type=str,
+    default="expand",
+    choices=[
+        "kernel", # for architecture except ResNet
+        "depth",
+        "expand",
+        "width", # only for ResNet
+    ],
+)
+parser.add_argument("--phase", type=int, default=2, choices=[1, 2])
+parser.add_argument("--resume", action="store_true")
+parser.add_argument("--model_name", type=str, default="MBV2", choices=["ResNet50", "MBV3", "ProxylessNASNet"])
+parser.add_argument("--dataset", type=str, default="cifar10", choices=["cifar10", "cifar100", "imagenet"])
+parser.add_argument("--robust_mode", type=bool, default=True)
+parser.add_argument("--epsilon", type=float, default=0.031)
+parser.add_argument("--num_steps", type=int, default=10)
+parser.add_argument("--step_size", type=float, default=0.0078) 
+parser.add_argument("--clip_min", type=int, default=0)
+parser.add_argument("--clip_max", type=int, default=1)
+parser.add_argument("--const_init", type=bool, default=False)
+parser.add_argument("--beta", type=float, default=6.0)
+parser.add_argument("--distance", type=str, default="l_inf",choices=["l_inf","l2"])
+parser.add_argument("--train_criterion", type=str, default="trades",choices=["trades","sat","mart","hat"])
+parser.add_argument("--test_criterion", type=str, default="ce",choices=["ce"])
+parser.add_argument("--kd_criterion", type=str, default="rslad",choices=["ard","rslad","adaad"])
+parser.add_argument("--attack_type", type=str, default="linf-pgd",choices=['fgsm', 'linf-pgd', 'fgm', 'l2-pgd', 'linf-df', 'l2-df', 'linf-apgd', 'l2-apgd','squar_attack','autoattack','apgd_ce'])
+args = parser.parse_args()
+if args.model_name == "ResNet50":
+    args.ks_list = "3"
+    if args.task == "width":
+        if args.robust_mode:
+            args.path = "exp/robust/WPS/"+ args.dataset + '/' + args.model_name +'/' + args.train_criterion +"/normal2width"
+        else:
+            args.path = "exp/WPS"+ args.dataset + '/' +args.model_name +'/' + args.train_criterion +"/normal2width"    
+        args.dynamic_batch_size = 1
+        args.n_epochs = 120
+        args.base_lr = 3e-2
+        args.warmup_epochs = 5
+        args.warmup_lr = -1
+        args.width_mult_list = "0.65,0.8,1.0"
+        args.expand_list = "0.35"
+        args.depth_list = "2"
+    elif args.task == "depth":
+        if args.robust_mode:
+            args.path = "exp/robust/WPS/"+ args.dataset + '/'  + args.model_name +'/' + args.train_criterion +"/width2width_depth/phase%d" % args.phase
+        else:
+            args.path = "exp/WPS/"+ args.dataset + '/'  + args.model_name +'/' + args.train_criterion +"/width2width_depth/phase%d" % args.phase     
+        args.dynamic_batch_size = 2
+        if args.phase == 1:
+            args.n_epochs = 25
+            args.base_lr = 2.5e-3
+            args.warmup_epochs = 0
+            args.warmup_lr = -1
+            args.width_mult_list = "0.65,0.8,1.0"
+            args.expand_list ="0.35"
+            args.depth_list = "1,2"
+        else:
+            args.n_epochs = 120
+            args.base_lr = 7.5e-3
+            args.warmup_epochs = 5
+            args.warmup_lr = -1
+            args.width_mult_list = "0.65,0.8,1.0"
+            args.expand_list = "0.35"
+            args.depth_list = "0,1,2"
+    elif args.task == "expand":
+        if args.robust_mode : 
+            args.path = "exp/robust/WPS/"+ args.dataset + '/' + args.model_name +'/' + args.train_criterion +"/width_depth2width_depth_width/phase%d" % args.phase 
+        else:
+            args.path = "exp/WPS/"+ args.dataset + '/' + args.model_name +'/' + args.train_criterion +"/width_depth2width_depth_width/phase%d" % args.phase    
+        args.dynamic_batch_size = 4
+        if args.phase == 1:
+            args.n_epochs = 25
+            args.base_lr = 2.5e-3
+            args.warmup_epochs = 0
+            args.warmup_lr = -1
+            args.width_mult_list = "0.65,0.8,1.0"
+            args.expand_list = "0.25,0.35"
+            args.depth_list = "0,1,2"
+        else:
+            args.n_epochs = 120
+            args.base_lr = 7.5e-3
+            args.warmup_epochs = 5
+            args.warmup_lr = -1
+            args.width_mult_list = "0.65,0.8,1.0"
+            args.expand_list = "0.2,0.25,0.35"
+            args.depth_list = "0,1,2"
+    else:
+        raise NotImplementedError
+else:
+    args.width_mult_list = "1.0"    
+    if args.task == "kernel":
+        if args.robust_mode:
+            args.path = "exp/robust/WPS/"+ args.dataset + '/' +  args.model_name +'/' + args.train_criterion +"/normal2kernel"
+        else:    
+            args.path = "exp/WPS/"+ args.dataset + '/' +  args.model_name +'/' + args.train_criterion +"/normal2kernel"
+        args.dynamic_batch_size = 1
+        args.n_epochs = 120
+        args.base_lr = 3e-2
+        args.warmup_epochs = 5
+        args.warmup_lr = -1
+        args.ks_list = "3,5,7"
+        args.expand_list = "6"
+        args.depth_list = "4"
+    elif args.task == "depth":
+        if args.robust_mode : 
+            args.path = "exp/robust/WPS/"+args.dataset + '/' + args.model_name +'/' + args.train_criterion +"/kernel2kernel_depth/phase%d" % args.phase
+        else:
+            args.path = "exp/WPS/"+args.dataset + '/' + args.model_name +'/' + args.train_criterion +"/kernel2kernel_depth/phase%d" % args.phase    
+        args.dynamic_batch_size = 2
+        if args.phase == 1:
+            args.n_epochs = 25
+            args.base_lr = 2.5e-3
+            args.warmup_epochs = 0
+            args.warmup_lr = -1
+            args.ks_list = "3,5,7"
+            args.expand_list = "6"
+            args.depth_list = "3,4"
+        else:
+            args.n_epochs = 120
+            args.base_lr = 7.5e-3
+            args.warmup_epochs = 5
+            args.warmup_lr = -1
+            args.ks_list = "3,5,7"
+            args.expand_list = "6"
+            args.depth_list = "2,3,4"
+    elif args.task == "expand":
+        if args.robust_mode:
+            args.path = "exp/robust/WPS/"+ args.dataset + '/' + args.model_name +'/' + args.train_criterion +"/kernel_depth2kernel_depth_width/phase%d" % args.phase
+        else:
+            args.path = "exp/WPS/"+ args.dataset + '/' + args.model_name +  '/' + args.train_criterion + "/kernel_depth2kernel_depth_width/phase%d" % args.phase    
+        args.dynamic_batch_size = 4
+        if args.phase == 1:
+            args.n_epochs = 25
+            args.base_lr = 2.5e-3
+            args.warmup_epochs = 0
+            args.warmup_lr = -1
+            args.ks_list = "3,5,7"
+            args.expand_list = "4,6"
+            args.depth_list = "2,3,4"
+        else:
+            args.n_epochs = 120
+            args.base_lr = 7.5e-3
+            args.warmup_epochs = 5
+            args.warmup_lr = -1
+            args.ks_list = "3,5,7"
+            args.expand_list = "3,4,6"
+            args.depth_list = "2,3,4"
+    else:
+        raise NotImplementedError
+args.manual_seed = 0
+
+args.lr_schedule_type = "cosine"
+
+args.base_batch_size = 64
+args.valid_size = 64
+
+args.opt_type = "sgd"
+args.momentum = 0.9
+args.no_nesterov = False
+args.weight_decay = 3e-5
+args.label_smoothing = 0.1
+args.no_decay_keys = "bn#bias"
+args.fp16_allreduce = False
+
+args.model_init = "he_fout"
+args.validation_frequency = 1
+args.print_frequency = 10
+
+args.n_worker = 8
+args.resize_scale = 0.08
+args.distort_color = "tf"
+if args.dataset == "imagenet":  
+    args.image_size = "128,160,192,224"
+else:
+    args.image_size = "32"    
+args.continuous_size = True
+args.not_sync_distributed_image_size = False
+
+args.bn_momentum = 0.1
+args.bn_eps = 1e-5
+args.dropout = 0.1
+args.base_stage_width = "google"
+
+
+args.dy_conv_scaling_mode = -1
+args.independent_distributed_sampling = False
+
+args.kd_ratio = 1.0
+args.kd_type = "ce"
+
+
+if __name__ == "__main__":
+    os.makedirs(args.path, exist_ok=True)
+
+    # Initialize Horovod
+    hvd.init()
+    # Pin GPU to be used to process local rank (one GPU per process)
+    torch.cuda.set_device(hvd.local_rank())
+    if args.robust_mode:
+        args.teacher_path = 'exp/robust/teacher/' + args.dataset + '/' +  args.model_name + '/' + args.train_criterion + "/checkpoint/model_best.pth.tar"
+    else:
+        args.teacher_path = 'exp/teacher/' + args.dataset + '/' +  args.model_name +'/' + args.train_criterion + "/checkpoint/model_best.pth.tar"
+    num_gpus = hvd.size()
+
+    torch.manual_seed(args.manual_seed)
+    torch.cuda.manual_seed_all(args.manual_seed)
+    np.random.seed(args.manual_seed)
+    random.seed(args.manual_seed)
+
+    # image size
+    args.image_size = [int(img_size) for img_size in args.image_size.split(",")]
+    if len(args.image_size) == 1:
+        args.image_size = args.image_size[0]
+    MyRandomResizedCrop.CONTINUOUS = args.continuous_size
+    MyRandomResizedCrop.SYNC_DISTRIBUTED = not args.not_sync_distributed_image_size
+
+    # build run config from args
+    args.lr_schedule_param = None
+    args.opt_param = {
+        "momentum": args.momentum,
+        "nesterov": not args.no_nesterov,
+    }
+    args.init_lr = args.base_lr * num_gpus  # linearly rescale the learning rate
+    if args.warmup_lr < 0:
+        args.warmup_lr = args.base_lr
+    args.train_batch_size = args.base_batch_size
+    args.test_batch_size = args.base_batch_size * 4
+    run_config = DistributedClassificationRunConfig(
+        **args.__dict__, num_replicas=num_gpus, rank=hvd.rank()
+    )
+
+    # print run config information
+    if hvd.rank() == 0:
+        print("Run config:")
+        for k, v in run_config.config.items():
+            print("\t%s: %s" % (k, v))
+
+    if args.dy_conv_scaling_mode == -1:
+        args.dy_conv_scaling_mode = None
+    DynamicSeparableConv2d.KERNEL_TRANSFORM_MODE = args.dy_conv_scaling_mode
+
+    # build net from args
+    args.width_mult_list = [
+        float(width_mult) for width_mult in args.width_mult_list.split(",")
+    ]
+    args.ks_list = [int(ks) for ks in args.ks_list.split(",")]
+    if args.model_name == "ResNet50":
+        args.expand_list = [float(e) for e in args.expand_list.split(",")]
+    else:
+        args.expand_list = [int(e) for e in args.expand_list.split(",")]    
+    args.depth_list = [int(d) for d in args.depth_list.split(",")]
+
+    args.width_mult_list = (
+        args.width_mult_list[0]
+        if len(args.width_mult_list) == 1
+        else args.width_mult_list
+    )
+
+    if args.model_name == "ResNet50":
+        if args.dataset == "cifar10" or args.dataset == "cifar100":
+            net = DYNResNets_Cifar( n_classes=run_config.data_provider.n_classes,
+                bn_param=(args.bn_momentum, args.bn_eps),
+                dropout_rate=args.dropout,
+                depth_list=args.depth_list,
+                expand_ratio_list=args.expand_list,
+                width_mult_list=args.width_mult_list,)
+        else:
+            net = DYNResNets( n_classes=run_config.data_provider.n_classes,
+                bn_param=(args.bn_momentum, args.bn_eps),
+                dropout_rate=args.dropout,
+                depth_list=args.depth_list,
+                expand_ratio_list=args.expand_list,
+                width_mult_list=args.width_mult_list,)   
+    elif args.model_name == "MBV3":  
+        if args.dataset == "cifar10" or args.dataset == "cifar100":  
+            net = DYNMobileNetV3_Cifar(n_classes=run_config.data_provider.n_classes,bn_param=(args.bn_momentum,args.bn_eps),
+                                dropout_rate= args.dropout, ks_list=args.ks_list , expand_ratio_list= args.expand_list , depth_list= args.depth_list,width_mult=args.width_mult_list)
+        else:
+            net = DYNMobileNetV3(n_classes=run_config.data_provider.n_classes,bn_param=(args.bn_momentum,args.bn_eps),
+                                dropout_rate= args.dropout, ks_list=args.ks_list , expand_ratio_list= args.expand_list , depth_list= args.depth_list,width_mult=args.width_mult_list)    
+    elif args.model_name == "ProxylessNASNet": 
+        if args.dataset == "cifar10" or args.dataset == "cifar100":     
+            net = DYNProxylessNASNets_Cifar(n_classes=run_config.data_provider.n_classes,bn_param=(args.bn_momentum,args.bn_eps),
+                                dropout_rate= args.dropout, ks_list=args.ks_list , expand_ratio_list= args.expand_list , depth_list= args.depth_list,width_mult=args.width_mult_list)
+        else:
+            net = DYNProxylessNASNets(n_classes=run_config.data_provider.n_classes,bn_param=(args.bn_momentum,args.bn_eps),
+                                dropout_rate= args.dropout, ks_list=args.ks_list , expand_ratio_list= args.expand_list , depth_list= args.depth_list,width_mult=args.width_mult_list)    
+    elif args.model_name == "MBV2": 
+        if args.dataset == "cifar10" or args.dataset == "cifar100":     
+            net = DYNProxylessNASNets_Cifar(n_classes=run_config.data_provider.n_classes,bn_param=(args.bn_momentum,args.bn_eps),
+                                dropout_rate= args.dropout, ks_list=args.ks_list , expand_ratio_list= args.expand_list , depth_list= args.depth_list,width_mult=args.width_mult_list,base_stage_width=args.base_stage_width)
+        else:
+            net = DYNProxylessNASNets(n_classes=run_config.data_provider.n_classes,bn_param=(args.bn_momentum,args.bn_eps),
+                                dropout_rate= args.dropout, ks_list=args.ks_list , expand_ratio_list= args.expand_list , depth_list= args.depth_list,width_mult=args.width_mult_list,base_stage_width=args.base_stage_width) 
+    else: 
+        raise NotImplementedError  
+    # teacher model
+    if args.kd_ratio > 0:
+        
+        if args.model_name =="ResNet50":
+            if args.dataset == "cifar10" or args.dataset == "cifar100":
+                args.teacher_model = DYNResNets_Cifar(
+                    n_classes=run_config.data_provider.n_classes,
+                    bn_param=(args.bn_momentum, args.bn_eps),
+                    dropout_rate=args.dropout,
+                    depth_list=[2],
+                    expand_ratio_list=[0.35],
+                    width_mult_list=[1.0],
+                )
+            else:
+                args.teacher_model = DYNResNets(
+                    n_classes=run_config.data_provider.n_classes,
+                    bn_param=(args.bn_momentum, args.bn_eps),
+                    dropout_rate=args.dropout,
+                    depth_list=[2],
+                    expand_ratio_list=[0.35],
+                    width_mult_list=[1.0],
+                )   
+        elif args.model_name =="MBV3":    
+            if args.dataset == "cifar10" or args.dataset == "cifar100":
+                args.teacher_model = DYNMobileNetV3_Cifar(
+                    n_classes=run_config.data_provider.n_classes,
+                    bn_param=(args.bn_momentum, args.bn_eps),
+                    dropout_rate=0,
+                    width_mult=1.0,
+                    ks_list=[7],
+                    expand_ratio_list=[6],
+                    depth_list=[4]
+                )
+            else:
+                args.teacher_model = DYNMobileNetV3(
+                    n_classes=run_config.data_provider.n_classes,
+                    bn_param=(args.bn_momentum, args.bn_eps),
+                    dropout_rate=0,
+                    width_mult=1.0,
+                    ks_list=[7],
+                    expand_ratio_list=[6],
+                    depth_list=[4]
+                )    
+        elif args.model_name == "ProxylessNASNet":
+            if args.dataset == "cifar10" or args.dataset == "cifar100":
+                args.teacher_model  = DYNProxylessNASNets_Cifar(n_classes=run_config.data_provider.n_classes,
+                    bn_param=(args.bn_momentum, args.bn_eps),
+                    dropout_rate=0,
+                    width_mult=1.0,
+                    ks_list=[7],
+                    expand_ratio_list=[6],
+                    depth_list=[4])   
+            else:
+                args.teacher_model  = DYNProxylessNASNets(n_classes=run_config.data_provider.n_classes,
+                    bn_param=(args.bn_momentum, args.bn_eps),
+                    dropout_rate=0,
+                    width_mult=1.0,
+                    ks_list=[7],
+                    expand_ratio_list=[6],
+                    depth_list=[4])      
+        elif args.model_name == "MBV2":
+            if args.dataset == "cifar10" or args.dataset == "cifar100":
+                args.teacher_model  = DYNProxylessNASNets_Cifar(n_classes=run_config.data_provider.n_classes,
+                    bn_param=(args.bn_momentum, args.bn_eps),
+                    dropout_rate=0,
+                    width_mult=1.0,
+                    ks_list=[7],
+                    expand_ratio_list=[6],
+                    depth_list=[4],base_stage_width=args.base_stage_width)   
+            else:
+                args.teacher_model  = DYNProxylessNASNets(n_classes=run_config.data_provider.n_classes,
+                    bn_param=(args.bn_momentum, args.bn_eps),
+                    dropout_rate=0,
+                    width_mult=1.0,
+                    ks_list=[7],
+                    expand_ratio_list=[6],
+                    depth_list=[4],base_stage_width=args.base_stage_width)     
+                        
+        args.teacher_model.cuda()
+    
+    """ Distributed RunManager """
+    # Horovod: (optional) compression algorithm.
+    compression = hvd.Compression.fp16 if args.fp16_allreduce else hvd.Compression.none
+    distributed_run_manager = DistributedRunManager(
+        args.path,
+        net,
+        run_config,
+        compression,
+        backward_steps=args.dynamic_batch_size,
+        is_root=(hvd.rank() == 0),
+    )
+    distributed_run_manager.save_config()
+    # hvd broadcast
+    distributed_run_manager.broadcast()
+
+    # load teacher net weights
+    if args.kd_ratio > 0:
+        load_models(
+            distributed_run_manager, args.teacher_model, model_path=args.teacher_path
+        )
+
+    # training
+    from proard.classification.elastic_nn.training.progressive_shrinking import (
+        validate,
+        train,
+    )
+    if args.model_name =="ResNet50":
+        validate_func_dict = {
+            "image_size_list": {224 if args.dataset == "imagenet" else 32}
+            if isinstance(args.image_size, int)
+            else sorted({160, 224}),
+            "width_mult_list": sorted({min(args.width_mult_list), max(args.width_mult_list)}),
+            "expand_ratio_list": sorted({min(args.expand_list), max(args.expand_list)}),
+            "depth_list": sorted({min(net.depth_list), max(net.depth_list)}),
+        }
+    else:
+        validate_func_dict = {
+            "image_size_list": {224 if args.dataset == "imagenet" else 32}
+            if isinstance(args.image_size, int)
+            else sorted({160, 224}),
+            "width_mult_list": [1.0],
+            "ks_list": sorted({min(args.ks_list), max(args.ks_list)}),
+            "expand_ratio_list": sorted({min(args.expand_list), max(args.expand_list)}),
+            "depth_list": sorted({min(net.depth_list), max(net.depth_list)}),
+        }  
+
+    if args.task == "width":
+        from proard.classification.elastic_nn.training.progressive_shrinking import (
+            train_elastic_width_mult,
+        )
+        if distributed_run_manager.start_epoch == 0:
+            if args.robust_mode:
+                args.dyn_checkpoint_path ='exp/robust/teacher/' +args.dataset + '/' +  args.model_name +'/' + args.train_criterion + "/checkpoint/model_best.pth.tar"
+            else:
+                args.dyn_checkpoint_path ='exp/teacher/' +args.dataset + '/' +  args.model_name +'/' + args.train_criterion + "/checkpoint/model_best.pth.tar"   
+            load_models(
+                distributed_run_manager,
+                distributed_run_manager.net,
+                args.dyn_checkpoint_path,
+            )
+            distributed_run_manager.write_log(
+                "%.3f\t%.3f\t%.3f\t%.3f\t%.3f\t%s"
+                % validate(distributed_run_manager, is_test=True, **validate_func_dict),
+                "valid",
+            )
+        else:
+            assert args.resume
+        train(distributed_run_manager,args,lambda _run_manager, epoch, is_test: validate(
+            _run_manager, epoch, is_test, **validate_func_dict
+        ),)      
+
+
+
+    elif args.task == "kernel":
+        validate_func_dict["ks_list"] = sorted(args.ks_list)
+        if distributed_run_manager.start_epoch == 0:
+            if args.robust_mode:
+                args.dyn_checkpoint_path ='exp/robust/teacher/' + args.dataset + '/' +  args.model_name +'/' + args.train_criterion + "/checkpoint/model_best.pth.tar"
+            else: 
+                args.dyn_checkpoint_path ='exp/teacher/' + args.dataset + '/' +  args.model_name +'/' + args.train_criterion + "/checkpoint/model_best.pth.tar"    
+            load_models(
+                distributed_run_manager,
+                distributed_run_manager.net,
+                args.dyn_checkpoint_path,
+            )
+            distributed_run_manager.write_log(
+               "%.3f\t%.3f\t%.3f\t%.3f\t%.3f\t%s"
+                % validate(distributed_run_manager, is_test=True, **validate_func_dict),
+                "valid",
+            )
+        else:
+            assert args.resume
+        train(
+            distributed_run_manager,
+            args,
+            lambda _run_manager, epoch, is_test: validate(
+                _run_manager, epoch, is_test, **validate_func_dict
+            ),
+        )
+    elif args.task == "depth":
+        from proard.classification.elastic_nn.training.progressive_shrinking import (
+            train_elastic_depth,
+        )
+        if args.robust_mode:
+            if args.model_name =="ResNet50":
+                if args.phase == 1:
+                    args.dyn_checkpoint_path =  "exp/robust/WPS/"+ args.dataset + '/'+ args.model_name +'/' + args.train_criterion +"/normal2width" +"/checkpoint/model_best.pth.tar"
+                else:
+                    args.dyn_checkpoint_path = "exp/robust/WPS/"+ args.dataset + '/' + args.model_name +'/' + args.train_criterion +"/width2width_depth/phase1" + "/checkpoint/model_best.pth.tar"
+            else:
+                if args.phase == 1:
+                    args.dyn_checkpoint_path =  "exp/robust/WPS/"+ args.dataset + '/' +  args.model_name +'/' + args.train_criterion +"/normal2kernel" +"/checkpoint/model_best.pth.tar"
+                else:
+                    args.dyn_checkpoint_path = "exp/robust/WPS/"+ args.dataset + '/' + args.model_name +'/' + args.train_criterion +"/kernel2kernel_depth/phase1" + "/checkpoint/model_best.pth.tar"     
+        else :
+            if args.model_name =="ResNet50":
+                if args.phase == 1:
+                    args.dyn_checkpoint_path =  "exp/WPS/"+ args.dataset + '/'+ args.model_name +'/' + args.train_criterion +"/normal2width" +"/checkpoint/model_best.pth.tar"
+                else:
+                    args.dyn_checkpoint_path = "exp/WPS/"+ args.dataset + '/' + args.model_name +'/' + args.train_criterion +"/width2width_depth/phase1" + "/checkpoint/model_best.pth.tar"
+            else:
+                if args.phase == 1:
+                    args.dyn_checkpoint_path =  "exp/WPS/"+ args.dataset + '/' +  args.model_name +'/' + args.train_criterion +"/normal2kernel" +"/checkpoint/model_best.pth.tar"
+                else:
+                    args.dyn_checkpoint_path = "exp/WPS/"+ args.dataset + '/' + args.model_name +'/' + args.train_criterion +"/kernel2kernel_depth/phase1" + "/checkpoint/model_best.pth.tar"                  
+        train(
+            distributed_run_manager,
+            args,
+            lambda _run_manager, epoch, is_test: validate(
+                _run_manager, epoch, is_test, **validate_func_dict
+            ),)
+    elif args.task == "expand":
+        from proard.classification.elastic_nn.training.progressive_shrinking import (
+            train_elastic_expand,
+        )
+        if args.robust_mode : 
+            if args.model_name =="ResNet50":
+                if args.phase == 1:
+                    args.dyn_checkpoint_path =  "exp/robust/WPS/"+ args.dataset + '/'+ args.model_name +'/' + args.train_criterion +"/width2width_depth/phase2" + "/checkpoint/model_best.pth.tar"
+                else:
+                    args.dyn_checkpoint_path = "exp/robust/WPS/"+ args.dataset + '/'+ args.model_name +'/' + args.train_criterion +"/width_depth2width_depth_width/phase1" + "/checkpoint/model_best.pth.tar"
+            else:
+                if args.phase == 1:
+                    args.dyn_checkpoint_path =  "exp/robust/WPS/"+ args.dataset + '/'+ args.model_name +'/' + args.train_criterion +"/kernel2kernel_depth/phase2" + "/checkpoint/model_best.pth.tar"  
+                else:
+                    args.dyn_checkpoint_path = "exp/robust/WPS/"+ args.dataset + '/'+ args.model_name +'/' + args.train_criterion +"/kernel_depth2kernel_depth_width/phase1" +  "/checkpoint/model_best.pth.tar" 
+        else:
+            if args.model_name =="ResNet50":
+                if args.phase == 1:
+                    args.dyn_checkpoint_path =  "exp/WPS/"+ args.dataset + '/'+ args.model_name +'/' + args.train_criterion +"/width2width_depth/phase2" + "/checkpoint/model_best.pth.tar"
+                else:
+                    args.dyn_checkpoint_path = "exp/WPS/"+ args.dataset + '/'+ args.model_name +'/' + args.train_criterion +"/width_depth2width_depth_width/phase1" + "/checkpoint/model_best.pth.tar"
+            else:
+                if args.phase == 1:
+                    args.dyn_checkpoint_path =  "exp/WPS/"+ args.dataset + '/'+ args.model_name +'/' + args.train_criterion +"/kernel2kernel_depth/phase2" + "/checkpoint/model_best.pth.tar"  
+                else:
+                    args.dyn_checkpoint_path = "exp/WPS/"+ args.dataset + '/'+ args.model_name +'/' + args.train_criterion +"/kernel_depth2kernel_depth_width/phase1" +  "/checkpoint/model_best.pth.tar" 
+
+        train(
+            distributed_run_manager,
+            args,
+            lambda _run_manager, epoch, is_test: validate(
+                _run_manager, epoch, is_test, **validate_func_dict
+            ),)
+    else:
+        raise NotImplementedError

train_teacher_net.py

@@ -0,0 +1,216 @@
+# Once for All: Train One Network and Specialize it for Efficient Deployment
+# Han Cai, Chuang Gan, Tianzhe Wang, Zhekai Zhang, Song Han
+# International Conference on Learning Representations (ICLR), 2020.
+
+import argparse
+import numpy as np
+import os
+import random
+# using for distributed training
+import horovod.torch as hvd
+import torch
+
+from proard.classification.elastic_nn.modules.dynamic_op import (
+    DynamicSeparableConv2d,
+)
+from proard.classification.elastic_nn.networks import DYNResNets,DYNMobileNetV3,DYNProxylessNASNets,DYNMobileNetV3_Cifar,DYNResNets_Cifar,DYNProxylessNASNets_Cifar
+from proard.classification.run_manager import DistributedClassificationRunConfig
+from proard.classification.networks import WideResNet
+from proard.classification.run_manager import DistributedRunManager
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--model_name", type=str, default="MBV2", choices=["ResNet50", "MBV3", "ProxylessNASNet","WideResNet","MBV2"])
+parser.add_argument("--teacher_model_name", type=str, default="WideResNet", choices=["WideResNet"])
+parser.add_argument("--dataset", type=str, default="cifar100", choices=["cifar10", "cifar100", "imagenet"])
+parser.add_argument("--robust_mode", type=bool, default=True)
+parser.add_argument("--epsilon", type=float, default=0.031)
+parser.add_argument("--num_steps", type=int, default=10)
+parser.add_argument("--step_size", type=float, default=0.0078) 
+parser.add_argument("--clip_min", type=int, default=0)
+parser.add_argument("--clip_max", type=int, default=1)
+parser.add_argument("--const_init", type=bool, default=False)
+parser.add_argument("--beta", type=float, default=6.0)
+parser.add_argument("--distance", type=str, default="l_inf",choices=["l_inf","l2"])
+parser.add_argument("--train_criterion", type=str, default="trades",choices=["trades","sat","mart","hat"])
+parser.add_argument("--test_criterion", type=str, default="ce",choices=["ce"])
+parser.add_argument("--kd_criterion", type=str, default="rslad",choices=["ard","rslad","adaad"])
+parser.add_argument("--attack_type", type=str, default="linf-pgd",choices=['fgsm', 'linf-pgd', 'fgm', 'l2-pgd', 'linf-df', 'l2-df', 'linf-apgd', 'l2-apgd','squar_attack','autoattack','apgd_ce'])
+
+
+args = parser.parse_args()
+if args.robust_mode:    
+    args.path = 'exp/robust/teacher/' + args.dataset + "/" +  args.model_name + '/' + args.train_criterion
+else:
+    args.path = 'exp/teacher/' + args.dataset + "/" +  args.model_name
+args.n_epochs = 120
+args.base_lr = 0.1
+args.warmup_epochs = 5
+args.warmup_lr = -1
+args.manual_seed = 0
+args.lr_schedule_type = "cosine"
+args.base_batch_size = 128
+args.valid_size = None
+args.opt_type = "sgd"
+args.momentum = 0.9
+args.no_nesterov = False
+args.weight_decay = 2e-4
+args.label_smoothing = 0.0
+args.no_decay_keys = "bn#bias"
+args.fp16_allreduce = False
+args.model_init = "he_fout"
+args.validation_frequency = 1
+args.print_frequency = 10
+args.n_worker = 32
+if args.dataset =="imagenet":
+    args.image_size = "224"
+else:
+    args.image_size = "32"     
+args.continuous_size = True
+args.not_sync_distributed_image_size = False
+args.bn_momentum = 0.1
+args.bn_eps = 1e-5
+args.dropout = 0.0
+args.base_stage_width = "google"
+###### Parameters for MBV3, ProxylessNet, and MBV2
+if args.model_name != "ResNet50":
+    args.ks_list = '7'
+    args.expand_list = '6'
+    args.depth_list = '4'
+    args.width_mult_list = "1.0"
+else: 
+    ###### Parameters for ResNet50
+    args.ks_list = "3"
+    args.expand_list = "0.35"
+    args.depth_list = "2"
+    args.width_mult_list = "1.0"
+########################################
+args.dy_conv_scaling_mode = 1
+args.independent_distributed_sampling = False
+args.kd_ratio = 0.0
+args.kd_type = "ce"
+args.dynamic_batch_size = 1
+args.num_gpus = 4
+if __name__ == "__main__":
+    os.makedirs(args.path, exist_ok=True)
+
+    # Initialize Horovod
+    hvd.init()
+    # Pin GPU to be used to process local rank (one GPU per process)
+    torch.cuda.set_device(hvd.local_rank())
+
+    num_gpus = hvd.size()
+    torch.manual_seed(args.manual_seed)
+    torch.cuda.manual_seed_all(args.manual_seed)
+    np.random.seed(args.manual_seed)
+    random.seed(args.manual_seed)
+
+    # image size
+    args.image_size = [int(img_size) for img_size in args.image_size.split(",")]
+    if len(args.image_size) == 1:
+        args.image_size = args.image_size[0]
+
+    # build run config from args
+    args.lr_schedule_param = None
+    args.opt_param = {
+        "momentum": args.momentum,
+        "nesterov": not args.no_nesterov,
+    }
+    args.init_lr = args.base_lr * num_gpus  # linearly rescale the learning rate
+    if args.warmup_lr < 0:
+        args.warmup_lr = args.base_lr
+    args.train_batch_size = args.base_batch_size
+    args.test_batch_size = args.base_batch_size 
+    print(args.__dict__)
+    run_config = DistributedClassificationRunConfig(
+        **args.__dict__,num_replicas=num_gpus, rank=hvd.rank()
+    )
+
+    # print run config information
+    if hvd.rank() == 0:
+        print("Run config:")
+        for k, v in run_config.config.items():
+            print("\t%s: %s" % (k, v))
+
+    if args.dy_conv_scaling_mode == -1:
+        args.dy_conv_scaling_mode = None
+    DynamicSeparableConv2d.KERNEL_TRANSFORM_MODE = args.dy_conv_scaling_mode
+
+    # build net from args
+    args.width_mult_list = [
+        float(width_mult) for width_mult in args.width_mult_list.split(",")
+    ]
+    args.ks_list = [int(ks) for ks in args.ks_list.split(",")]
+    args.expand_list = [float(e) for e in args.expand_list.split(",")]
+    args.depth_list = [int(d) for d in args.depth_list.split(",")]
+
+    args.width_mult_list = (
+        args.width_mult_list[0]
+        if len(args.width_mult_list) == 1
+        else args.width_mult_list
+    )
+    if args.model_name == "ResNet50":
+        if args.dataset == "cifar10" or args.dataset == "cifar100":
+            # net = ResNet50_Cifar(n_classes=run_config.data_provider.n_classes)
+            net = DYNResNets_Cifar( n_classes=run_config.data_provider.n_classes,
+                bn_param=(args.bn_momentum, args.bn_eps),
+                dropout_rate=args.dropout,
+                depth_list=args.depth_list,
+                expand_ratio_list=args.expand_list,
+                width_mult_list=args.width_mult_list,)
+        else:
+            net = DYNResNets( n_classes=run_config.data_provider.n_classes,
+                bn_param=(args.bn_momentum, args.bn_eps),
+                dropout_rate=args.dropout,
+                depth_list=args.depth_list,
+                expand_ratio_list=args.expand_list,
+                width_mult_list=args.width_mult_list,)   
+    elif args.model_name == "MBV3":  
+        if args.dataset == "cifar10" or args.dataset == "cifar100":  
+            net = DYNMobileNetV3_Cifar(n_classes=run_config.data_provider.n_classes,bn_param=(args.bn_momentum,args.bn_eps),
+                                dropout_rate= args.dropout, ks_list=args.ks_list , expand_ratio_list= args.expand_list , depth_list= args.depth_list)
+        else:
+            net = DYNMobileNetV3(n_classes=run_config.data_provider.n_classes,bn_param=(args.bn_momentum,args.bn_eps),
+                                dropout_rate= args.dropout, ks_list=args.ks_list , expand_ratio_list= args.expand_list , depth_list= args.depth_list)    
+    elif args.model_name == "ProxylessNASNet": 
+        if args.dataset == "cifar10" or args.dataset == "cifar100":     
+            net = DYNProxylessNASNets_Cifar(n_classes=run_config.data_provider.n_classes,bn_param=(args.bn_momentum,args.bn_eps),
+                                dropout_rate= args.dropout, ks_list=args.ks_list , expand_ratio_list= args.expand_list , depth_list= args.depth_list)
+        else:
+            net = DYNProxylessNASNets(n_classes=run_config.data_provider.n_classes,bn_param=(args.bn_momentum,args.bn_eps),
+                                dropout_rate= args.dropout, ks_list=args.ks_list , expand_ratio_list= args.expand_list , depth_list= args.depth_list) 
+             
+    elif args.model_name == "MBV2": 
+        if args.dataset == "cifar10" or args.dataset == "cifar100":     
+            net = DYNProxylessNASNets_Cifar(n_classes=run_config.data_provider.n_classes,bn_param=(args.bn_momentum,args.bn_eps),base_stage_width=args.base_stage_width,
+                                dropout_rate= args.dropout, ks_list=args.ks_list , expand_ratio_list= args.expand_list , depth_list= args.depth_list)
+        else:
+            net = DYNProxylessNASNets(n_classes=run_config.data_provider.n_classes,bn_param=(args.bn_momentum,args.bn_eps),base_stage_width=args.base_stage_width,
+                                dropout_rate= args.dropout, ks_list=args.ks_list , expand_ratio_list= args.expand_list , depth_list= args.depth_list) 
+    else: 
+        raise NotImplementedError                    
+    if args.teacher_model_name == "WideResNet":  
+        if args.dataset == "cifar10" or args.dataset == "cifar100":  
+            net = WideResNet(num_classes=run_config.data_provider.n_classes)
+        else:
+           raise NotImplementedError 
+    else: 
+        raise NotImplementedError  
+    args.teacher_model = None #'exp/teacher/' + args.dataset + "/" +  "WideResNet"
+    
+    """ Distributed RunManager """
+    #Horovod: (optional) compression algorithm.
+    compression = hvd.Compression.fp16 if args.fp16_allreduce else hvd.Compression.none
+    distributed_run_manager = DistributedRunManager(
+        args.path,
+        net,
+        run_config,
+        compression,
+        backward_steps=args.dynamic_batch_size,
+        is_root=(hvd.rank() == 0),
+    )
+    distributed_run_manager.save_config()
+    distributed_run_manager.broadcast()
+
+   
+    distributed_run_manager.train(args)