diff --git a/args_helper.py b/args_helper.py index 94370efa..f4e76804 100644 --- a/args_helper.py +++ b/args_helper.py @@ -282,8 +282,6 @@ def parse_arguments(self, jupyter_mode=False): default=1.0, help="portion of additional width compared with original width" ) - - parser.add_argument( "--hidden-size", type=int, @@ -899,14 +897,12 @@ def parse_arguments(self, jupyter_mode=False): default=False, help="Only run sanity checks on the files in specific directory or subdirectories" ) - parser.add_argument( "--invert-sanity-check", action="store_true", default=False, help="Enable this to run the inverted sanity check (for HC)" ) - parser.add_argument( "--sanity-folder", default=None, @@ -914,13 +910,18 @@ def parse_arguments(self, jupyter_mode=False): metavar="PATH", help="directory(s) to access for only sanity check", ) - parser.add_argument( "--sr-version", default=1, type=int, help="smart ratio version number (1, 2, ...)", ) + parser.add_argument( + "--only-finetune", + action="store_true", + default=False, + help="Enable this to skip pruning and jump to finetune. Typically paired with --resume" + ) parser.add_argument( "--use-full-data", action="store_true", diff --git a/configs/hypercube/resnet18/imagenet/resnet18_sparsity_10.yml b/configs/hypercube/resnet18/imagenet/resnet18_sparsity_10.yml new file mode 100644 index 00000000..b044f695 --- /dev/null +++ b/configs/hypercube/resnet18/imagenet/resnet18_sparsity_10.yml @@ -0,0 +1,66 @@ +subfolder: ffcv_imagenet_resnet18 +# trial_num: 1 +#lam_finetune_loss: 1 +#num_step_finetune: 5 + +# Hypercube optimization +algo: 'hc_iter' +iter_period: 5 + +# Architecture +arch: ResNet18 + +# ===== Dataset ===== # +dataset: FfcvImageNet +name: resnet18_ffcv +data: /workspace/ffcv-imagenet/data/ + +# ===== Learning Rate Policy ======== # +optimizer: sgd +lr: 0.1 #0.01 +lr_policy: cosine_lr #constant_lr #multistep_lr +fine_tune_lr: 0.01 +fine_tune_lr_policy: multistep_lr + +# ===== Network training config ===== # +epochs: 50 +wd: 0.0 +momentum: 0.9 +batch_size: 128 +mixed_precision: True + +# ===== Sparsity =========== # +conv_type: SubnetConv +bn_type: NonAffineBatchNorm +freeze_weights: True +prune_type: BottomK +# enter target sparsity here +target_sparsity: 10 +# decide if you want to "unflag" +unflag_before_finetune: False +init: signed_constant +score_init: unif #skew #half #bimodal #skew # bern +scale_fan: False #True + +# ===== Rounding ===== # +round: naive +noise: True +noise_ratio: 0 + +# ===== Quantization ===== # +hc_quantized: True +quantize_threshold: 0.5 + +# ===== Regularization ===== # +regularization: L2 +lmbda: 0.000001 # 1e-4 #0.00005 # 5e-5 + +# ===== Hardware setup ===== # +workers: 4 +gpu: 0 + +# ===== Checkpointing ===== # +checkpoint_at_prune: False + +# ==== sanity check ==== # +skip_sanity_checks: True diff --git a/configs/hypercube/resnet20/sreeniva_resnet20_1_4_best.yml b/configs/hypercube/resnet20/sreeniva_resnet20_1_4_best.yml new file mode 100644 index 00000000..5385628f --- /dev/null +++ b/configs/hypercube/resnet20/sreeniva_resnet20_1_4_best.yml @@ -0,0 +1,64 @@ +# subfolder: sreeniva_hc_1_44_best +# trial_num: 1 +#lam_finetune_loss: 1 +#num_step_finetune: 5 + +# Hypercube optimization +algo: 'hc_iter' +iter_period: 5 + +# Architecture +arch: resnet20 + +# ===== Dataset ===== # +dataset: CIFAR10 +name: resnet20_quantized_iter_hc + +# ===== Learning Rate Policy ======== # +optimizer: sgd +lr: 0.1 #0.01 +lr_policy: cosine_lr #constant_lr #multistep_lr +fine_tune_lr: 0.01 +fine_tune_lr_policy: multistep_lr + +# ===== Network training config ===== # +epochs: 150 +wd: 0.0 +momentum: 0.9 +batch_size: 128 + +# ===== Sparsity =========== # +conv_type: SubnetConv +bn_type: NonAffineBatchNorm +freeze_weights: True +prune_type: BottomK +# enter target sparsity here +target_sparsity: 1.44 +# decide if you want to "unflag" +unflag_before_finetune: False +init: signed_constant +score_init: unif #skew #half #bimodal #skew # bern +scale_fan: False #True + +# ===== Rounding ===== # +round: naive +noise: True +noise_ratio: 0 + +# ===== Quantization ===== # +hc_quantized: True +quantize_threshold: 0.5 + +# ===== Regularization ===== # +regularization: L2 +lmbda: 0.0001 # 1e-4 #0.00005 # 5e-5 + +# ===== Hardware setup ===== # +workers: 4 +gpu: 2 + +# ===== Checkpointing ===== # +checkpoint_at_prune: True + +# ==== sanity check ==== # +skip_sanity_checks: True diff --git a/configs/hypercube/resnet50/ffcv_imagenet/resnet50_sparsity_10_finetune_only.yml b/configs/hypercube/resnet50/ffcv_imagenet/resnet50_sparsity_10_finetune_only.yml new file mode 100644 index 00000000..b3effa24 --- /dev/null +++ b/configs/hypercube/resnet50/ffcv_imagenet/resnet50_sparsity_10_finetune_only.yml @@ -0,0 +1,68 @@ +subfolder: ffcv_imagenet_resnet50_finetune_only +# trial_num: 1 +#lam_finetune_loss: 1 +#num_step_finetune: 5 + +# Hypercube optimization +algo: 'hc_iter' +iter_period: 5 + +# Architecture +arch: ResNet50 + +# ===== Dataset ===== # +dataset: FfcvImageNet +name: resnet18_ffcv +data: /workspace/ffcv-imagenet/data/ + +# ===== Learning Rate Policy ======== # +optimizer: sgd +lr: 0.4 #0.01 +lr_policy: cosine_lr #constant_lr #multistep_lr +fine_tune_lr: 0.01 +fine_tune_lr_policy: multistep_lr + +# ===== Network training config ===== # +epochs: 88 +wd: 0.0 +momentum: 0.9 +batch_size: 128 +mixed_precision: True +only_finetune: True +resume: results/resnet50_ffcv_bad_finetune/model_before_finetune.pth + +# ===== Sparsity =========== # +conv_type: SubnetConv +bn_type: NonAffineBatchNorm +freeze_weights: True +prune_type: BottomK +# enter target sparsity here +target_sparsity: 10 +# decide if you want to "unflag" +unflag_before_finetune: False +init: signed_constant +score_init: unif #skew #half #bimodal #skew # bern +scale_fan: False #True + +# ===== Rounding ===== # +round: naive +noise: True +noise_ratio: 0 + +# ===== Quantization ===== # +hc_quantized: True +quantize_threshold: 0.5 + +# ===== Regularization ===== # +regularization: L2 +lmbda: 0.000001 # 1e-4 #0.00005 # 5e-5 + +# ===== Hardware setup ===== # +workers: 6 +gpu: 0 + +# ===== Checkpointing ===== # +checkpoint_at_prune: False + +# ==== sanity check ==== # +skip_sanity_checks: True diff --git a/configs/hypercube/resnet50/ffcv_imagenet/resnet50_sparsity_20.yml b/configs/hypercube/resnet50/ffcv_imagenet/resnet50_sparsity_20.yml new file mode 100644 index 00000000..564b8a45 --- /dev/null +++ b/configs/hypercube/resnet50/ffcv_imagenet/resnet50_sparsity_20.yml @@ -0,0 +1,67 @@ +subfolder: ffcv_imagenet_resnet50_sp20 +# trial_num: 1 +#lam_finetune_loss: 1 +#num_step_finetune: 5 + +# Hypercube optimization +algo: 'hc_iter' +iter_period: 5 + +# Architecture +arch: ResNet50 + +# ===== Dataset ===== # +dataset: FfcvImageNet +name: resnet50_ffcv_imagenet +data: /workspace/ffcv-imagenet/data/ + +# ===== Learning Rate Policy ======== # +optimizer: sgd +lr: 0.1 #0.01 +lr_policy: cosine_lr #constant_lr #multistep_lr +fine_tune_lr: 0.01 +fine_tune_lr_policy: multistep_lr + +# ===== Network training config ===== # +epochs: 88 +wd: 0.0 +momentum: 0.9 +batch_size: 256 +mixed_precision: True + +# ===== Sparsity =========== # +conv_type: SubnetConv +bn_type: NonAffineBatchNorm +freeze_weights: True +prune_type: BottomK +# enter target sparsity here +target_sparsity: 20 +# decide if you want to "unflag" +unflag_before_finetune: False +init: signed_constant +score_init: unif #skew #half #bimodal #skew # bern +scale_fan: False #True +use_full_data: True + +# ===== Rounding ===== # +round: naive +noise: True +noise_ratio: 0 + +# ===== Quantization ===== # +hc_quantized: True +quantize_threshold: 0.5 + +# ===== Regularization ===== # +regularization: L2 +lmbda: 0.000001 # 1e-4 #0.00005 # 5e-5 + +# ===== Hardware setup ===== # +workers: 4 +gpu: 1 + +# ===== Checkpointing ===== # +checkpoint_at_prune: False + +# ==== sanity check ==== # +skip_sanity_checks: False diff --git a/configs/hypercube/resnet50/ffcv_imagenet/resnet50_sparsity_5.yml b/configs/hypercube/resnet50/ffcv_imagenet/resnet50_sparsity_5.yml new file mode 100644 index 00000000..b89f53cf --- /dev/null +++ b/configs/hypercube/resnet50/ffcv_imagenet/resnet50_sparsity_5.yml @@ -0,0 +1,66 @@ +subfolder: ffcv_imagenet_resnet50 +# trial_num: 1 +#lam_finetune_loss: 1 +#num_step_finetune: 5 + +# Hypercube optimization +algo: 'hc_iter' +iter_period: 5 + +# Architecture +arch: ResNet50 + +# ===== Dataset ===== # +dataset: FfcvImageNet +name: resnet18_ffcv +data: /workspace/ffcv-imagenet/data/ + +# ===== Learning Rate Policy ======== # +optimizer: sgd +lr: 0.4 #0.01 +lr_policy: cosine_lr #constant_lr #multistep_lr +fine_tune_lr: 0.001 +fine_tune_lr_policy: multistep_lr + +# ===== Network training config ===== # +epochs: 88 +wd: 0.0 +momentum: 0.9 +batch_size: 256 +mixed_precision: True + +# ===== Sparsity =========== # +conv_type: SubnetConv +bn_type: NonAffineBatchNorm +freeze_weights: True +prune_type: BottomK +# enter target sparsity here +target_sparsity: 5 +# decide if you want to "unflag" +unflag_before_finetune: False +init: signed_constant +score_init: unif #skew #half #bimodal #skew # bern +scale_fan: False #True + +# ===== Rounding ===== # +round: naive +noise: True +noise_ratio: 0 + +# ===== Quantization ===== # +hc_quantized: True +quantize_threshold: 0.5 + +# ===== Regularization ===== # +regularization: L2 +lmbda: 0.0000001 # 1e-4 #0.00005 # 5e-5 + +# ===== Hardware setup ===== # +workers: 8 +gpu: 0 + +# ===== Checkpointing ===== # +checkpoint_at_prune: False + +# ==== sanity check ==== # +skip_sanity_checks: True diff --git a/configs/hypercube/resnet50/ffcv_imagenet/resnet50_sparsity_50.yml b/configs/hypercube/resnet50/ffcv_imagenet/resnet50_sparsity_50.yml new file mode 100644 index 00000000..97eeb818 --- /dev/null +++ b/configs/hypercube/resnet50/ffcv_imagenet/resnet50_sparsity_50.yml @@ -0,0 +1,66 @@ +subfolder: ffcv_imagenet_resnet50_sp50 +# trial_num: 1 +#lam_finetune_loss: 1 +#num_step_finetune: 5 + +# Hypercube optimization +algo: 'hc_iter' +iter_period: 5 + +# Architecture +arch: ResNet50 + +# ===== Dataset ===== # +dataset: FfcvImageNet +name: resnet18_ffcv +data: /workspace/ffcv-imagenet/data/ + +# ===== Learning Rate Policy ======== # +optimizer: sgd +lr: 0.1 #0.01 +lr_policy: cosine_lr #constant_lr #multistep_lr +fine_tune_lr: 0.001 +fine_tune_lr_policy: multistep_lr + +# ===== Network training config ===== # +epochs: 88 +wd: 0.0 +momentum: 0.9 +batch_size: 200 +mixed_precision: True + +# ===== Sparsity =========== # +conv_type: SubnetConv +bn_type: NonAffineBatchNorm +freeze_weights: True +prune_type: BottomK +# enter target sparsity here +target_sparsity: 50 +# decide if you want to "unflag" +unflag_before_finetune: False +init: signed_constant +score_init: unif #skew #half #bimodal #skew # bern +scale_fan: False #True + +# ===== Rounding ===== # +round: naive +noise: True +noise_ratio: 0 + +# ===== Quantization ===== # +hc_quantized: True +quantize_threshold: 0.5 + +# ===== Regularization ===== # +regularization: L2 +lmbda: 0 # 1e-4 #0.00005 # 5e-5 + +# ===== Hardware setup ===== # +workers: 8 +gpu: 1 + +# ===== Checkpointing ===== # +checkpoint_at_prune: False + +# ==== sanity check ==== # +skip_sanity_checks: True diff --git a/configs/hypercube/resnet50/ffcv_imagenet/resnet50_sparsity_5_v2.yml b/configs/hypercube/resnet50/ffcv_imagenet/resnet50_sparsity_5_v2.yml new file mode 100644 index 00000000..f5c62b1b --- /dev/null +++ b/configs/hypercube/resnet50/ffcv_imagenet/resnet50_sparsity_5_v2.yml @@ -0,0 +1,66 @@ +subfolder: ffcv_imagenet_resnet50 +# trial_num: 1 +#lam_finetune_loss: 1 +#num_step_finetune: 5 + +# Hypercube optimization +algo: 'hc_iter' +iter_period: 5 + +# Architecture +arch: ResNet50 + +# ===== Dataset ===== # +dataset: FfcvImageNet +name: resnet18_ffcv +data: /workspace/ffcv-imagenet/data/ + +# ===== Learning Rate Policy ======== # +optimizer: sgd +lr: 0.1 #0.01 +lr_policy: cosine_lr #constant_lr #multistep_lr +fine_tune_lr: 0.001 +fine_tune_lr_policy: multistep_lr + +# ===== Network training config ===== # +epochs: 88 +wd: 0.0 +momentum: 0.9 +batch_size: 256 +mixed_precision: True + +# ===== Sparsity =========== # +conv_type: SubnetConv +bn_type: NonAffineBatchNorm +freeze_weights: True +prune_type: BottomK +# enter target sparsity here +target_sparsity: 5 +# decide if you want to "unflag" +unflag_before_finetune: False +init: signed_constant +score_init: unif #skew #half #bimodal #skew # bern +scale_fan: False #True + +# ===== Rounding ===== # +round: naive +noise: True +noise_ratio: 0 + +# ===== Quantization ===== # +hc_quantized: True +quantize_threshold: 0.5 + +# ===== Regularization ===== # +regularization: L2 +lmbda: 0.0000001 # 1e-4 #0.00005 # 5e-5 + +# ===== Hardware setup ===== # +workers: 8 +gpu: 0 + +# ===== Checkpointing ===== # +checkpoint_at_prune: False + +# ==== sanity check ==== # +skip_sanity_checks: True diff --git a/configs/hypercube/resnet50/imagenet/resnet50_sparsity_5.yml b/configs/hypercube/resnet50/imagenet/resnet50_sparsity_5.yml new file mode 100644 index 00000000..3576abcf --- /dev/null +++ b/configs/hypercube/resnet50/imagenet/resnet50_sparsity_5.yml @@ -0,0 +1,66 @@ +subfolder: regular_imagenet_resnet50 +# trial_num: 1 +#lam_finetune_loss: 1 +#num_step_finetune: 5 + +# Hypercube optimization +algo: 'hc_iter' +iter_period: 5 + +# Architecture +arch: ResNet50 + +# ===== Dataset ===== # +dataset: ImageNet +name: resnet50_imagenet +data: /data/imagenet/ + +# ===== Learning Rate Policy ======== # +optimizer: sgd +lr: 0.4 #0.01 +lr_policy: cosine_lr #constant_lr #multistep_lr +fine_tune_lr: 0.001 +fine_tune_lr_policy: multistep_lr + +# ===== Network training config ===== # +epochs: 10 +wd: 0.0 +momentum: 0.9 +batch_size: 64 +mixed_precision: True + +# ===== Sparsity =========== # +conv_type: SubnetConv +bn_type: NonAffineBatchNorm +freeze_weights: True +prune_type: BottomK +# enter target sparsity here +target_sparsity: 5 +# decide if you want to "unflag" +unflag_before_finetune: False +init: signed_constant +score_init: unif #skew #half #bimodal #skew # bern +scale_fan: False #True + +# ===== Rounding ===== # +round: naive +noise: True +noise_ratio: 0 + +# ===== Quantization ===== # +hc_quantized: True +quantize_threshold: 0.5 + +# ===== Regularization ===== # +regularization: L2 +lmbda: 0.0000001 # 1e-4 #0.00005 # 5e-5 + +# ===== Hardware setup ===== # +workers: 4 +gpu: 1 + +# ===== Checkpointing ===== # +checkpoint_at_prune: False + +# ==== sanity check ==== # +skip_sanity_checks: True diff --git a/data/__init__.py b/data/__init__.py index a70f282b..596b67d6 100644 --- a/data/__init__.py +++ b/data/__init__.py @@ -3,3 +3,4 @@ from data.tinyimagenet import TinyImageNet from data.mnist import MNIST from data.bigcifar import BigCIFAR10 +from data.ffcv_imagenet import FfcvImageNet diff --git a/data/ffcv_imagenet.py b/data/ffcv_imagenet.py new file mode 100644 index 00000000..665e8a74 --- /dev/null +++ b/data/ffcv_imagenet.py @@ -0,0 +1,177 @@ +""" + script to create ffcv_imagenet train loader +""" + +import os +import torch +import torchvision +from torchvision import datasets, transforms +import random +from torch.utils.data.sampler import SubsetRandomSampler +from args_helper import parser_args +import numpy as np +from pathlib import Path +from typing import List + +from ffcv.pipeline.operation import Operation +from ffcv.loader import Loader, OrderOption +from ffcv.transforms import ToTensor, ToDevice, Squeeze, NormalizeImage, \ + RandomHorizontalFlip, ToTorchImage +from ffcv.fields.rgb_image import CenterCropRGBImageDecoder, \ + RandomResizedCropRGBImageDecoder +from ffcv.fields.basics import IntDecoder + +torch.backends.cudnn.benchmark = True +torch.autograd.profiler.emit_nvtx(False) +torch.autograd.profiler.profile(False) + +import torch.multiprocessing +torch.multiprocessing.set_sharing_strategy("file_system") + +from torch.utils.data import random_split + +class FfcvImageNet: + def __init__(self, args): + super(FfcvImageNet, self).__init__() + + # data_root = os.path.join(parser_args.data, "imagenet") + # put ffcv path here + data_root = parser_args.data + + use_cuda = torch.cuda.is_available() + + # Data loading code + dataset = os.path.join(data_root, "train_500_0.50_90.ffcv") + val_dataset = os.path.join(data_root, "val_500_0.50_90.ffcv") + + # Data loading code + kwargs = {"num_workers": 6, "in_memory": 1, + "distributed": False, "resolution": 256} + + self.IMAGENET_MEAN = np.array([0.485, 0.456, 0.406]) * 255 + self.IMAGENET_STD = np.array([0.229, 0.224, 0.225]) * 255 + self.DEFAULT_CROP_RATIO = 224/256 + + if parser_args.use_full_data: + train_dataset = dataset + # use_full_data => we are not tuning hyperparameters + actual_val_dataset = val_dataset + else: + # this will break right now. But where is the test set even? + val_size = 5000 + train_size = len(dataset) - val_size + train_dataset, actual_val_dataset = random_split(dataset, [train_size, val_size]) + + + self.train_loader = self.create_train_loader(train_dataset, + kwargs['num_workers'], + parser_args.batch_size, + kwargs['distributed'], + kwargs['in_memory'] + ) + self.actual_val_loader = self.create_train_loader(actual_val_dataset, + kwargs['num_workers'], + parser_args.batch_size, + kwargs['distributed'], + kwargs['in_memory'] + ) + self.val_loader = self.create_val_loader(val_dataset, + kwargs['num_workers'], + parser_args.batch_size, + kwargs['resolution'], + kwargs['distributed'], + ) + # madry does this but I don't think we need to + # self.model, self.scaler = self.create_model_and_scaler() + + def create_train_loader(self, train_dataset, num_workers, batch_size, + distributed, in_memory): + train_path = Path(train_dataset) + assert train_path.is_file() + + res = self.get_resolution(epoch=0) + self.decoder = RandomResizedCropRGBImageDecoder((res, res)) + image_pipeline: List[Operation] = [ + self.decoder, + RandomHorizontalFlip(), + ToTensor(), + ToDevice(torch.device(parser_args.gpu), non_blocking=True), + ToTorchImage(), + NormalizeImage(self.IMAGENET_MEAN, self.IMAGENET_STD, np.float16) + ] + + label_pipeline: List[Operation] = [ + IntDecoder(), + ToTensor(), + Squeeze(), + ToDevice(torch.device(parser_args.gpu), non_blocking=True) + ] + + order = OrderOption.RANDOM if distributed else OrderOption.QUASI_RANDOM + loader = Loader(train_dataset, + batch_size=batch_size, + num_workers=num_workers, + order=order, + os_cache=in_memory, + drop_last=True, + pipelines={ + 'image': image_pipeline, + 'label': label_pipeline + }, + distributed=distributed) + + return loader + + + def create_val_loader(self, val_dataset, num_workers, batch_size, + resolution, distributed): + val_path = Path(val_dataset) + assert val_path.is_file() + res_tuple = (resolution, resolution) + cropper = CenterCropRGBImageDecoder(res_tuple, ratio=self.DEFAULT_CROP_RATIO) + image_pipeline = [ + cropper, + ToTensor(), + ToDevice(torch.device(parser_args.gpu), non_blocking=True), + ToTorchImage(), + NormalizeImage(self.IMAGENET_MEAN, self.IMAGENET_STD, np.float16) + ] + + label_pipeline = [ + IntDecoder(), + ToTensor(), + Squeeze(), + ToDevice(torch.device(parser_args.gpu), + non_blocking=True) + ] + + loader = Loader(val_dataset, + batch_size=batch_size, + num_workers=num_workers, + order=OrderOption.SEQUENTIAL, + drop_last=False, + pipelines={ + 'image': image_pipeline, + 'label': label_pipeline + }, + distributed=distributed) + return loader + + def get_resolution(self, epoch=0, min_res=160, max_res=192, end_ramp=76, start_ramp=65): + # this seems to be a hack to get good accuracy and is only between epochs + # 65 and 76. So, for now just return max_res. always. + + return max_res + + assert min_res <= max_res + + if epoch <= start_ramp: + return min_res + + if epoch >= end_ramp: + return max_res + + # otherwise, linearly interpolate to the nearest multiple of 32 + interp = np.interp([epoch], [start_ramp, end_ramp], [min_res, max_res]) + final_res = int(np.round(interp[0] / 32)) * 32 + return final_res diff --git a/ffcv_installation.sh b/ffcv_installation.sh new file mode 100644 index 00000000..9a6af4b2 --- /dev/null +++ b/ffcv_installation.sh @@ -0,0 +1,41 @@ +# set up ffcv script + +docker run --gpus all --ipc=host --name kartik_imagenet -it -v /hdd1/ILSVRC2012:/data/imagenet -v /hdd3/ksreenivasan:/workspace nvcr.io/nvidia/pytorch:21.12-py3 + +sudo docker pull nvcr.io/nvidia/pytorch:22.01-py3 + +# clone repo +git clone git@github.com:libffcv/ffcv-imagenet.git + +conda create -y -n ffcv python=3.9 cupy pkg-config compilers libjpeg-turbo opencv pytorch torchvision cudatoolkit=11.3 numba -c pytorch -c conda-forge +conda activate ffcv +pip install ffcv + +# to fix dependencies for ffcv +apt-get install ffmpeg libsm6 libxext6 -y + +# install dependencies +cd ffcv-imagenet +pip install -r requirements.txt + +# set environment variables +export IMAGENET_DIR=/data/imagenet +export WRITE_DIR=/workspace/ffcv-imagenet/data +mkdir data + +# create ffcv data + +# Serialize images with: +# - 500px side length maximum +# - 50% JPEG encoded, 90% raw pixel values +# - quality=90 JPEGs +./write_imagenet.sh 500 0.50 90 + + +python train_imagenet.py --config-file rn18_configs/rn18_88_epochs.yaml \ + --data.train_dataset=/workspace/ffcv-imagenet/data/train_500_0.50_90.ffcv \ + --data.val_dataset=/workspace/ffcv-imagenet/data/val_500_0.50_90.ffcv \ + --data.num_workers=12 --data.in_memory=1 \ + --logging.folder=/workspace/ffcv-imagenet/logs + +# start exec: 5:24pm (with batch size 512) \ No newline at end of file diff --git a/imagenet_exec.sh b/imagenet_exec.sh new file mode 100644 index 00000000..f8c08079 --- /dev/null +++ b/imagenet_exec.sh @@ -0,0 +1,31 @@ +#### ResNet-18 + +# Running trials in parallel +# NOTE: make sure to delete/comment subfolder from the config file or else it may not work +:< "$log_root$trial$log_end" 2>&1 & + + python main.py \ + --config "$conf_file" \ + --trial-num $trial \ + --invert-sanity-check \ + --subfolder "invert_$subfolder_root$trial" > "invert_$log_root$trial$log_end" 2>&1 & +done + +BLOCK + +conf_file="configs/hypercube/resnet50/ffcv_imagenet/resnet50_sparsity_50.yml" +log_root="resnet50_ffcv_sp50" +log_end="_log" +python main.py \ + --config "$conf_file" > "$log_root$log_end" 2>&1 & diff --git a/main.py b/main.py index 42dc3829..bf5ee4e0 100644 --- a/main.py +++ b/main.py @@ -135,6 +135,10 @@ def main_worker(gpu, ngpus_per_node): # Start training for epoch in range(parser_args.start_epoch, parser_args.epochs): + if parser_args.only_finetune: + print("Skipping pruning and going straight to finetune!!!") + break + if parser_args.multiprocessing_distributed: data.train_loader.sampler.set_epoch(epoch) #lr_policy(epoch, iteration=None) @@ -279,6 +283,10 @@ def main_worker(gpu, ngpus_per_node): print("Writing results into: {}".format(results_filename)) results_df.to_csv(results_filename, index=False) + if parser_args.resume: + print("Loading checkpoint before finetune") + best_acc1 = resume(parser_args, model, optimizer) + # save checkpoint before fine-tuning torch.save(model.state_dict(), result_root + 'model_before_finetune.pth') diff --git a/main_utils.py b/main_utils.py index 0ab9a80b..60f32cc3 100644 --- a/main_utils.py +++ b/main_utils.py @@ -997,6 +997,13 @@ def get_model(parser_args): if parser_args.freeze_weights: freeze_model_weights(model) + # TODO: hacky. fix this. @sreeniva + #if parser_args.arch == "ResNet18": + # # make half precision + # model = model.half() + # for name, params in model.named_parameters(): + # params.data = params.half() + return model diff --git a/models/__init__.py b/models/__init__.py index 216a0bfb..1c769ae3 100644 --- a/models/__init__.py +++ b/models/__init__.py @@ -7,7 +7,9 @@ #### TODO: delete below ones (merge with above code) from models.resnet_cifar import cResNet18, cResNet50 -from models.resnet_tiny import TinyResNet18 +from models.resnet_tiny import TinyResNet18 + +# NOTE: ResNet18 is for imagenet - rename __all__ = [ "tinyvgg16", diff --git a/sanity_check_prospr.py b/sanity_check_prospr.py new file mode 100644 index 00000000..9d2e6caa --- /dev/null +++ b/sanity_check_prospr.py @@ -0,0 +1,161 @@ +# trying to test prospr using their codebase + +import importlib +import re + +from train import * +import argparse +import random +from pathlib import Path + +import utils +from cli import parse_args +import prospr.utils +import random +import numpy as np +import os +import copy + +from models.resnet20 import _weights_init + +WEIGHT_REINIT_SANITY = False +MASK_SHUFFLE_SANITY = True + +args = parse_args() +hparams = utils.Hyperparameters(**vars(args)) + +log_dir = utils.create_logdir("sreeniva_debug_prospr") + +utils.set_seed(hparams.seed, hparams.allow_nondeterminism) + +train_data, _, test_data, _ = dataloader_factory( + hparams.dataset, hparams.batch_size +) + +model = model_factory(hparams.model, hparams.dataset, hparams.no_model_patching) +model, masks = get_pruned_model(model, hparams) + +model_ckpt = torch.load("/workspace/pruning_is_enough/prospr_ckpts/sp95/trained_model.pt") +mask_ckpt = torch.load("/workspace/pruning_is_enough/prospr_ckpts/sp95/pruning_keep_mask.pt") +model.load_state_dict(model_ckpt) +masks = mask_ckpt + +filter_fn = pruning_filter_factory(10, hparams.structured_pruning) +structured = False + +pruned_model = prospr.utils.apply_masks_with_hooks(model, masks, structured, filter_fn) + +# from now onwards, only refer to pruned_model!!! +optimizer, lr_scheduler = get_optimizer(pruned_model, hparams) +pruned_model = pruned_model.cuda() + +test_loss, test_acc1, test_acc5 = evaluate(pruned_model, test_data) +print("Test Accuracy: {}%".format(test_acc1*100)) +orig_test_acc = test_acc1*100 + +weight_only = [] +for name, param in pruned_model.named_parameters(): + if "weight" in name and "bn" not in name: + weight_only.append((name, param)) + +# manually apply mask to model as sanity check +idx = 0 +for name, param in pruned_model.named_parameters(): + if "weight" in name and "bn" not in name: + param.data *= masks[idx] + idx += 1 + +# this should be the same as before +test_loss, test_acc1, test_acc5 = evaluate(pruned_model, test_data) +print("Test Accuracy: {}%".format(test_acc1*100)) + + +# Run sanity checks +def set_seed(seed): + random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + np.random.seed(seed) + os.environ['PYTHONHASHSEED'] = str(seed) + # making sure GPU runs are deterministic even if they are slower + torch.backends.cudnn.deterministic = True + # this causes the code to vary across runs. I don't want that for now. + # torch.backends.cudnn.benchmark = True + print("Seeded everything: {}".format(seed)) + +set_seed(96) + +bkp_pruned_model = copy.deepcopy(pruned_model) + +# weight reinit sanity check +if WEIGHT_REINIT_SANITY: + pruned_model = pruned_model.apply(_weights_init) + print("Running Weight Reinit Sanity Check") +else: + shuffled_masks = [] + for mask in masks: + idx = torch.randperm(mask.nelement()) + shuffled_masks.append(mask.view(-1)[idx].view(mask.size())) + print("Running Mask Shuffle Sanity Check") + masks = shuffled_masks + +pruned_model = prospr.utils.apply_masks_with_hooks(pruned_model, masks, structured, filter_fn) +# manually apply mask to model as sanity check +idx = 0 +for name, param in pruned_model.named_parameters(): + if "weight" in name and "bn" not in name: + print("Pruning Layer: {}".format(name)) + print("Before: {}".format(torch.norm(param))) + param.data *= masks[idx] + print("After: {}".format(torch.norm(param))) + idx += 1 + +# compare norms to verify that reinit has really happened. +print("Conv1.weight norm: Before={} | After={}".format(torch.norm(bkp_pruned_model.conv1.weight), torch.norm(pruned_model.conv1.weight))) + +# this should be the same as before +test_loss, test_acc1, test_acc5 = evaluate(pruned_model, test_data) +print("Test Accuracy after flipping things (should be terrible): {}%".format(test_acc1*100)) + +# train for 200 epochs and see what happens +optimizer, lr_scheduler = get_optimizer(pruned_model, hparams) +for epoch in range(1, 200 + 1): + avg_train_loss, epoch_time = train_one_epoch(pruned_model, train_data, optimizer) + test_loss, test_acc1, test_acc5 = evaluate(pruned_model, test_data) + + print( + f"📸 Epoch {epoch} (finished in {epoch_time})\n", + f"\tTrain loss:\t{avg_train_loss:.4f}\n", + f"\tTest loss:\t{test_loss:.4f}\n", + f"\tTest acc:\t{test_acc1:.4f}\n", + f"\tTest top-5 acc:\t{test_acc5:.4f}", + ) + + lr_scheduler.step() + +print( + "✅ Training finished\n", + f"\tFinal test acc: {test_acc1}\n", + f"\tFinal test acc@5: {test_acc5}", +) + +# sanity check by applying the mask and testing accuracy +# manually apply mask to model as sanity check +idx = 0 +for name, param in pruned_model.named_parameters(): + if "weight" in name and "bn" not in name: + print("Pruning Layer: {}".format(name)) + print("Before: {}".format(torch.norm(param))) + param.data *= masks[idx] + print("After: {}".format(torch.norm(param))) + idx += 1 + +# this should be the same as before +test_loss, test_acc1, test_acc5 = evaluate(pruned_model, test_data) +print("Test Accuracy: {}%".format(test_acc1*100)) +final_test_acc = test_acc1*100 + +print("Sanity check complete!") +print("Accuracy before reinit: {} | Accuracy after weight reinit: {}".format(orig_test_acc, final_test_acc)) + diff --git a/trainers/default.py b/trainers/default.py index 6597353f..70ff3304 100644 --- a/trainers/default.py +++ b/trainers/default.py @@ -7,13 +7,13 @@ from utils.eval_utils import accuracy from utils.logging import AverageMeter, ProgressMeter from utils.net_utils import get_regularization_loss, prune, get_layers +from torch.cuda.amp import autocast from torch import optim __all__ = ["train", "validate", "modifier"] - def train(train_loader, model, criterion, optimizer, epoch, args, writer, scaler=None): batch_time = AverageMeter("Time", ":6.3f") data_time = AverageMeter("Data", ":6.3f") @@ -57,17 +57,13 @@ def train(train_loader, model, criterion, optimizer, epoch, args, writer, scaler scores.data = torch.clamp(scores.data, 0.0, 1.0) # compute output - if scaler is None: + with torch.cuda.amp.autocast(enabled=args.mixed_precision): # mixed precision output = model(images) loss = criterion(output, target) - else: - with torch.cuda.amp.autocast(enabled=True): # mixed precision - output = model(images) - loss = criterion(output, target) if args.lam_finetune_loss > 0: raise NotImplementedError # please check finetune_loss repo - + regularization_loss = torch.tensor(0) if args.regularization: regularization_loss =\ @@ -76,7 +72,9 @@ def train(train_loader, model, criterion, optimizer, epoch, args, writer, scaler alpha_prime=args.alpha_prime) #print('regularization_loss: ', regularization_loss) - loss += regularization_loss + + with torch.cuda.amp.autocast(enabled=args.mixed_precision): # mixed precision + loss += regularization_loss # measure accuracy and record loss acc1, acc5, acc10 = accuracy(output, target, topk=(1, 5, 10)) @@ -91,6 +89,8 @@ def train(train_loader, model, criterion, optimizer, epoch, args, writer, scaler loss.backward() optimizer.step() else: + with torch.cuda.amp.autocast(enabled=args.mixed_precision): # mixed precision + output = model(images) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update() @@ -145,9 +145,9 @@ def validate(val_loader, model, criterion, args, writer, epoch): #print(images.shape, target.shape) # compute output - output = model(images) - - loss = criterion(output, target) + with torch.cuda.amp.autocast(enabled=args.mixed_precision): # mixed precision + output = model(images) + loss = criterion(output, target) # measure accuracy and record loss acc1, acc5, acc10 = accuracy(output, target, topk=(1, 5, 10)) diff --git a/utils/conv_type.py b/utils/conv_type.py index c235d14b..99109708 100644 --- a/utils/conv_type.py +++ b/utils/conv_type.py @@ -71,11 +71,11 @@ def forward(ctx, scores, bias_scores, k, scores_prune_threshold=-np.inf, bias_sc # NOTE: doing this EP style where the scores are unchanged, but mask is computed # can also try a variant where we actually round the scores if parser_args.bottom_k_on_forward: - out = torch.gt(scores, torch.ones_like(scores)*scores_prune_threshold).float() - bias_out = torch.gt(bias_scores, torch.ones_like(bias_scores)*bias_scores_prune_threshold).float() + out = torch.gt(scores, torch.ones_like(scores)*scores_prune_threshold).float().half() + bias_out = torch.gt(bias_scores, torch.ones_like(bias_scores)*bias_scores_prune_threshold).float().half() else: - out = torch.gt(scores, torch.ones_like(scores)*parser_args.quantize_threshold).float() - bias_out = torch.gt(bias_scores, torch.ones_like(bias_scores)*parser_args.quantize_threshold).float() + out = torch.gt(scores, torch.ones_like(scores)*parser_args.quantize_threshold).float().half() + bias_out = torch.gt(bias_scores, torch.ones_like(bias_scores)*parser_args.quantize_threshold).float().half() else: print("INVALID PRUNING ALGO") @@ -171,11 +171,11 @@ def forward(self, x): if parser_args.hc_quantized: subnet, bias_subnet = GetSubnet.apply(self.scores, self.bias_scores, parser_args.prune_rate) - subnet = subnet * self.flag.data.float() - bias_subnet = subnet * self.bias_flag.data.float() + subnet = subnet * self.flag.data.float().half() + bias_subnet = subnet * self.bias_flag.data.float().half() else: - subnet = self.scores * self.flag.data.float() - bias_subnet = self.bias_scores * self.bias_flag.data.float() + subnet = self.scores * self.flag.data.float().half() + bias_subnet = self.bias_scores * self.bias_flag.data.float().half() elif parser_args.algo in ['imp']: # no STE, no subnet. Mask is handled outside pass @@ -195,7 +195,8 @@ def forward(self, x): b = self.bias * bias_subnet else: b = self.bias - + # TODO: wrap the ".half()" in something + # print("dtype: {}".format(w.dtype)) x = F.conv2d( x, w, b, self.stride, self.padding, self.dilation, self.groups ) diff --git a/utils/net_utils.py b/utils/net_utils.py index 46462981..3aeedd29 100644 --- a/utils/net_utils.py +++ b/utils/net_utils.py @@ -126,6 +126,18 @@ def get_layers(arch='Conv4', model=None): linear_layers.append(layer.mlp.fc1) linear_layers.append(layer.mlp.fc2) # linear_layers.append(model.decoder) + + elif arch == 'ResNet18': + conv_layers = [model.conv1] + for layer in [model.layer1, model.layer2, model.layer3, model.layer4]: + for basic_block_id in [0, 1]: + conv_layers.append(layer[basic_block_id].conv1) + conv_layers.append(layer[basic_block_id].conv2) + linear_layers = [model.fc] + + else: + print("ERROR: get_layers() not implemented for model {}".format(arch)) + return -1 return (conv_layers, linear_layers) diff --git a/utils/utils.py b/utils/utils.py index 1db4f656..2ed92dba 100644 --- a/utils/utils.py +++ b/utils/utils.py @@ -25,9 +25,9 @@ def set_seed(seed): np.random.seed(seed) os.environ['PYTHONHASHSEED'] = str(seed) # making sure GPU runs are deterministic even if they are slower - torch.backends.cudnn.deterministic = True + torch.backends.cudnn.deterministic = False # this causes the code to vary across runs. I don't want that for now. - # torch.backends.cudnn.benchmark = True + torch.backends.cudnn.benchmark = True print("Seeded everything: {}".format(seed))