Why Do We Need Weight Decay in Modern Deep Learning? [arXiv, Oct 2023]
# config for training GPT-2 (124M) down to very nice loss of ~2.85 on 1 node of 8X A100 40GB
# launch as the following (e.g. in a screen session) and wait ~5 days:
# $ torchrun --standalone --nproc_per_node=8 train.py config/train_gpt2.py

wandb_log = True
wandb_project = 'llm'
wandb_run_name = 'owt'

# these make the total batch size be ~0.5M
# 12 batch size * 1024 block size * 5 gradaccum * 8 GPUs = 491,520
batch_size = 12
gradient_accumulation_steps = 40

# this makes total number of tokens be 300B
max_iters = 600000

# model
n_layer = 4  # default: 12
n_head = 4  # default: 12
n_embd = 192  # default: 768
block_size = 256  # default: 1024

# eval stuff
eval_interval = 200
eval_examples = 1000
log_interval = 10

# LR stuff
warmup_iters = 400 # default: 2000

# weight decay
weight_decay = 0.0

large_language_models/config/train_gpt2_micro.py
tml-epfl/why-weight-decay
# config for training GPT-2 (124M) down to very nice loss of ~2.85 on 1 node of 8X A100 40GB
# launch as the following (e.g. in a screen session) and wait ~5 days:
# $ torchrun --standalone --nproc_per_node=8 train.py config/train_gpt2.py

wandb_log = True
wandb_project = 'llm'
wandb_run_name='owt'

# these make the total batch size be ~0.5M
# 12 batch size * 1024 block size * 5 gradaccum * 8 GPUs = 491,520
batch_size = 12
gradient_accumulation_steps = 40

# this makes total number of tokens be 300B
max_iters = 600000

# model
n_layer = 12  # default: 12
n_head = 12  # default: 12
n_embd = 768  # default: 768
block_size = 1024  # default: 1024

# eval stuff
eval_interval = 200
eval_examples = 1000
log_interval = 10

# LR stuff
warmup_iters = 400 # default: 2000

# weight decay
weight_decay = 0.0

large_language_models/config/train_gpt2_small.py
tml-epfl/why-weight-decay
# config for training GPT-2 (124M) down to very nice loss of ~2.85 on 1 node of 8X A100 40GB
# launch as the following (e.g. in a screen session) and wait ~5 days:
# $ torchrun --standalone --nproc_per_node=8 train.py config/train_gpt2.py

wandb_log = True
wandb_project = 'llm'
wandb_run_name='owt'

# these make the total batch size be ~0.5M
# 12 batch size * 1024 block size * 5 gradaccum * 8 GPUs = 491,520
batch_size = 12
gradient_accumulation_steps = 40

# this makes total number of tokens be 300B
max_iters = 600000

# model
n_layer = 12  # default: 12
n_head = 12  # default: 12
n_embd = 768  # default: 768
block_size = 256  # default: 1024

# eval stuff
eval_interval = 200
eval_examples = 1000
log_interval = 10

# LR stuff
warmup_iters = 400 # default: 2000

# weight decay
weight_decay = 0.0

large_language_models/config/train_gpt2_small_block256.py
tml-epfl/why-weight-decay
# train a miniature character-level shakespeare model
# good for debugging and playing on macbooks and such

eval_interval = 250 # keep frequent because we'll overfit
eval_iters = 200
log_interval = 10 # don't print too too often

# we expect to overfit on this small dataset, so only save when val improves
always_save_checkpoint = False

wandb_log = True # override via command line if you like
wandb_project = 'llm'
wandb_run_name = 'shakespeare'  # the WD value will be appended later

dataset = 'shakespeare_char'
gradient_accumulation_steps = 1
batch_size = 64
block_size = 256 # context of up to 256 previous characters

# baby GPT model :)
n_layer = 6
n_head = 6
n_embd = 384
weight_decay = 0.0  # default: 0.1
dropout = 0.0  # default: 0.2

learning_rate = 1e-3 # with baby networks can afford to go a bit higher
max_iters = 10000
lr_decay_iters = max_iters # make equal to max_iters usually
# min_lr = 1e-4 # learning_rate / 10 usually
beta2 = 0.99 # make a bit bigger because number of tokens per iter is small

warmup_iters = 100 # not super necessary potentially

# on macbook also add
# device = 'cpu'  # run on cpu only
# compile = False # do not torch compile the model

large_language_models/config/train_shakespeare_char.py
tml-epfl/why-weight-decay
"""
Poor Man's Configurator. Probably a terrible idea. Example usage:
$ python train.py config/override_file.py --batch_size=32
this will first run config/override_file.py, then override batch_size to 32

The code in this file will be run as follows from e.g. train.py:
>>> exec(open('configurator.py').read())

So it's not a Python module, it's just shuttling this code away from train.py
The code in this script then overrides the globals()

I know people are not going to love this, I just really dislike configuration
complexity and having to prepend config. to every single variable. If someone
comes up with a better simple Python solution I am all ears.
"""

import sys
from ast import literal_eval

for arg in sys.argv[1:]:
    if '=' not in arg:
        # assume it's the name of a config file
        assert not arg.startswith('--')
        config_file = arg
        print(f"Overriding config with {config_file}:")
        with open(config_file) as f:
            print(f.read())
        exec(open(config_file).read())
    else:
        # assume it's a --key=value argument
        assert arg.startswith('--')
        key = arg.split('=')[0]
        val = '='.join(arg.split('=')[1:])
        key = key[2:]
        if key in globals():
            try:
                # attempt to eval it it (e.g. if bool, number, or etc)
                attempt = literal_eval(val)
            except (SyntaxError, ValueError):
                # if that goes wrong, just use the string
                attempt = val
            # ensure the types match ok; upd: breaks on weight_decay=0
            # assert type(attempt) == type(globals()[key])
            # cross fingers
            print(f"Overriding: {key} = {attempt}")
            globals()[key] = attempt
        else:
            raise ValueError(f"Unknown config key: {key}")

large_language_models/configurator.py
tml-epfl/why-weight-decay
"""
Full definition of a GPT Language Model, all of it in this single file.
References:
1) the official GPT-2 TensorFlow implementation released by OpenAI:
https://github.com/openai/gpt-2/blob/master/src/model.py
2) huggingface/transformers PyTorch implementation:
https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt2/modeling_gpt2.py
"""

import math
import inspect
from dataclasses import dataclass

import numpy as np
import torch
import torch.nn as nn
from torch.nn import functional as F

class LayerNorm(nn.Module):
    """ LayerNorm but with an optional bias. PyTorch doesn't support simply bias=False """

    def __init__(self, ndim, bias, learnable_ln):
        super().__init__()
        self.ndim = ndim
        self.weight = nn.Parameter(torch.ones(ndim)) if learnable_ln else None
        self.bias = nn.Parameter(torch.zeros(ndim)) if bias else None

    def forward(self, input):
        # return F.layer_norm(input, self.weight.shape, self.weight, self.bias, 1e-5)
        return F.layer_norm(input, [self.ndim], self.weight, self.bias, 1e-5)

class CausalSelfAttention(nn.Module):

    def __init__(self, config):
        super().__init__()
        assert config.n_embd % config.n_head == 0
        # key, query, value projections for all heads, but in a batch
        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)
        # output projection
        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)
        # regularization
        self.attn_dropout = nn.Dropout(config.dropout)
        self.resid_dropout = nn.Dropout(config.dropout)
        self.n_head = config.n_head
        self.n_embd = config.n_embd
        self.dropout = config.dropout
        # flash attention make GPU go brrrrr but support is only in PyTorch >= 2.0
        self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention')
        if not self.flash:
            print("WARNING: using slow attention. Flash Attention requires PyTorch >= 2.0")
            # causal mask to ensure that attention is only applied to the left in the input sequence
            self.register_buffer("bias", torch.tril(torch.ones(config.block_size, config.block_size))
                                        .view(1, 1, config.block_size, config.block_size))

    def forward(self, x):
        B, T, C = x.size() # batch size, sequence length, embedding dimensionality (n_embd)

        # calculate query, key, values for all heads in batch and move head forward to be the batch dim
        q, k, v  = self.c_attn(x).split(self.n_embd, dim=2)
        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)

        # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)
        if self.flash:
            # efficient attention using Flash Attention CUDA kernels
            y = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=self.dropout if self.training else 0, is_causal=True)
        else:
            # manual implementation of attention
            att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
            att = att.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf'))
            att = F.softmax(att, dim=-1)
            att = self.attn_dropout(att)
            y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
        y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side

        # output projection
        y = self.resid_dropout(self.c_proj(y))
        return y

class MLP(nn.Module):

    def __init__(self, config):
        super().__init__()
        self.c_fc    = nn.Linear(config.n_embd, 4 * config.n_embd, bias=config.bias)
        self.gelu    = nn.GELU()
        self.c_proj  = nn.Linear(4 * config.n_embd, config.n_embd, bias=config.bias)
        self.dropout = nn.Dropout(config.dropout)

    def forward(self, x):
        x = self.c_fc(x)
        x = self.gelu(x)
        x = self.c_proj(x)
        x = self.dropout(x)
        return x

class Block(nn.Module):

    def __init__(self, config):
        super().__init__()
        self.ln_1 = LayerNorm(config.n_embd, bias=config.bias, learnable_ln=config.learnable_ln)
        self.attn = CausalSelfAttention(config)
        self.ln_2 = LayerNorm(config.n_embd, bias=config.bias, learnable_ln=config.learnable_ln)
        self.mlp = MLP(config)

    def forward(self, x):
        x = x + self.attn(self.ln_1(x))
        x = x + self.mlp(self.ln_2(x))
        return x

@dataclass
class GPTConfig:
    block_size: int = 1024
    vocab_size: int = 50304 # GPT-2 vocab_size of 50257, padded up to nearest multiple of 64 for efficiency
    n_layer: int = 12
    n_head: int = 12
    n_embd: int = 768
    dropout: float = 0.0
    bias: bool = True # True: bias in Linears and LayerNorms, like GPT-2. False: a bit better and faster
    learnable_ln: bool = True
    weight_tying: bool = True
    custom_init: bool = False
    init_scale: float = 0.02

class GPT(nn.Module):

    def __init__(self, config):
        super().__init__()
        assert config.vocab_size is not None
        assert config.block_size is not None
        self.config = config

        self.transformer = nn.ModuleDict(dict(
            wte = nn.Embedding(config.vocab_size, config.n_embd),
            wpe = nn.Embedding(config.block_size, config.n_embd),
            drop = nn.Dropout(config.dropout),
            h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
            ln_f = LayerNorm(config.n_embd, bias=config.bias, learnable_ln=config.learnable_ln),
        ))
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
        # with weight tying when using torch.compile() some warnings get generated:
        # "UserWarning: functional_call was passed multiple values for tied weights.
        # This behavior is deprecated and will be an error in future versions"
        # not 100% sure what this is, so far seems to be harmless. TODO investigate
        if config.weight_tying:
            self.transformer.wte.weight = self.lm_head.weight # https://paperswithcode.com/method/weight-tying

        # init all weights
        self.apply(self._init_weights)
        # apply special scaled init to the residual projections, per GPT-2 paper
        for pn, p in self.named_parameters():
            if pn.endswith('c_proj.weight'):
                torch.nn.init.normal_(p, mean=0.0, std=self.config.init_scale/math.sqrt(2 * config.n_layer))
        
        if config.custom_init:
            # note: it's not clear how to initialize the first layer since it's tied with the last one
            for pn, p in self.named_parameters():
                if len(p.shape) > 1:
                    nn.init.kaiming_normal_(p, mode='fan_in', nonlinearity='linear')
                    if 'wte.weight' in pn or 'wpe.weight' in pn:
                        p.data = p.data / 2**0.5
                    elif 'attn' in pn or 'mlp' in pn:
                        p.data = p.data / p.shape[-1]**0.5
                    # p.data = p.data / (2 * config.n_layer)
                    print(pn, p.shape, torch.mean((p**2).sum(1)).item())

        # report number of parameters
        print("number of parameters: %.2fM" % (self.get_num_params()/1e6,))

    def get_num_params(self, non_embedding=True):
        """
        Return the number of parameters in the model.
        For non-embedding count (default), the position embeddings get subtracted.
        The token embeddings would too, except due to the parameter sharing these
        params are actually used as weights in the final layer, so we include them.
        """
        n_params = sum(p.numel() for p in self.parameters())
        if non_embedding:
            n_params -= self.transformer.wpe.weight.numel()
        return n_params

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            torch.nn.init.normal_(module.weight, mean=0.0, std=self.config.init_scale)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0.0, std=self.config.init_scale)

    def forward(self, idx, targets=None):
        device = idx.device
        b, t = idx.size()
        assert t <= self.config.block_size, f"Cannot forward sequence of length {t}, block size is only {self.config.block_size}"
        pos = torch.arange(0, t, dtype=torch.long, device=device) # shape (t)

        # forward the GPT model itself
        tok_emb = self.transformer.wte(idx) # token embeddings of shape (b, t, n_embd)
        pos_emb = self.transformer.wpe(pos) # position embeddings of shape (t, n_embd)
        x = self.transformer.drop(tok_emb + pos_emb)
        for block in self.transformer.h:
            # print(x.shape, (x**2).sum(-1).mean([0, 1]))
            x = block(x)
        # print(x.shape, (x**2).sum(-1).mean([0, 1]))
        # import ipdb;ipdb.set_trace()
        x = self.transformer.ln_f(x)

        if targets is not None:
            # if we are given some desired targets also calculate the loss
            logits = self.lm_head(x)
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)
        else:
            # inference-time mini-optimization: only forward the lm_head on the very last position
            logits = self.lm_head(x[:, [-1], :]) # note: using list [-1] to preserve the time dim
            loss = None

        return logits, loss

    def crop_block_size(self, block_size):
        # model surgery to decrease the block size if necessary
        # e.g. we may load the GPT2 pretrained model checkpoint (block size 1024)
        # but want to use a smaller block size for some smaller, simpler model
        assert block_size <= self.config.block_size
        self.config.block_size = block_size
        self.transformer.wpe.weight = nn.Parameter(self.transformer.wpe.weight[:block_size])
        for block in self.transformer.h:
            if hasattr(block.attn, 'bias'):
                block.attn.bias = block.attn.bias[:,:,:block_size,:block_size]

    @classmethod
    def from_pretrained(cls, model_type, override_args=None):
        assert model_type in {'gpt2', 'gpt2-medium', 'gpt2-large', 'gpt2-xl'}
        override_args = override_args or {} # default to empty dict
        # only dropout can be overridden see more notes below
        assert all(k == 'dropout' for k in override_args)
        from transformers import GPT2LMHeadModel
        print("loading weights from pretrained gpt: %s" % model_type)

        # n_layer, n_head and n_embd are determined from model_type
        config_args = {
            'gpt2':         dict(n_layer=12, n_head=12, n_embd=768),  # 124M params
            'gpt2-medium':  dict(n_layer=24, n_head=16, n_embd=1024), # 350M params
            'gpt2-large':   dict(n_layer=36, n_head=20, n_embd=1280), # 774M params
            'gpt2-xl':      dict(n_layer=48, n_head=25, n_embd=1600), # 1558M params
        }[model_type]
        print("forcing vocab_size=50257, block_size=1024, bias=True")
        config_args['vocab_size'] = 50257 # always 50257 for GPT model checkpoints
        config_args['block_size'] = 1024 # always 1024 for GPT model checkpoints
        config_args['bias'] = True # always True for GPT model checkpoints
        # we can override the dropout rate, if desired
        if 'dropout' in override_args:
            print(f"overriding dropout rate to {override_args['dropout']}")
            config_args['dropout'] = override_args['dropout']
        # create a from-scratch initialized minGPT model
        config = GPTConfig(**config_args)
        model = GPT(config)
        sd = model.state_dict()
        sd_keys = sd.keys()
        sd_keys = [k for k in sd_keys if not k.endswith('.attn.bias')] # discard this mask / buffer, not a param

        # init a huggingface/transformers model
        model_hf = GPT2LMHeadModel.from_pretrained(model_type)
        sd_hf = model_hf.state_dict()

        # copy while ensuring all of the parameters are aligned and match in names and shapes
        sd_keys_hf = sd_hf.keys()
        sd_keys_hf = [k for k in sd_keys_hf if not k.endswith('.attn.masked_bias')] # ignore these, just a buffer
        sd_keys_hf = [k for k in sd_keys_hf if not k.endswith('.attn.bias')] # same, just the mask (buffer)
        transposed = ['attn.c_attn.weight', 'attn.c_proj.weight', 'mlp.c_fc.weight', 'mlp.c_proj.weight']
        # basically the openai checkpoints use a "Conv1D" module, but we only want to use a vanilla Linear
        # this means that we have to transpose these weights when we import them
        assert len(sd_keys_hf) == len(sd_keys), f"mismatched keys: {len(sd_keys_hf)} != {len(sd_keys)}"
        for k in sd_keys_hf:
            if any(k.endswith(w) for w in transposed):
                # special treatment for the Conv1D weights we need to transpose
                assert sd_hf[k].shape[::-1] == sd[k].shape
                with torch.no_grad():
                    sd[k].copy_(sd_hf[k].t())
            else:
                # vanilla copy over the other parameters
                assert sd_hf[k].shape == sd[k].shape
                with torch.no_grad():
                    sd[k].copy_(sd_hf[k])

        return model

    def configure_optimizers(self, opt_type, weight_decay, learning_rate, betas, device_type, wd_substrings_include, dampening):
        # start with all of the candidate parameters
        param_dict = {pn: p for pn, p in self.named_parameters()}
        # filter out those that do not require grad
        param_dict = {pn: p for pn, p in param_dict.items() if p.requires_grad}
        # create optim groups. Any parameters that is 2D will be weight decayed, otherwise no.
        # i.e. all weight tensors in matmuls + embeddings decay, all biases and layernorms don't.
        decay_params, nodecay_params = [], []
        for n, p in param_dict.items():
            if any([s in n for s in wd_substrings_include]):
                decay_params.append(p)
            else:
                nodecay_params.append(p)
        optim_groups = [
            {'params': decay_params, 'weight_decay': weight_decay},
            {'params': nodecay_params, 'weight_decay': 0.0}
        ]
        num_decay_params = sum(p.numel() for p in decay_params)
        num_nodecay_params = sum(p.numel() for p in nodecay_params)
        print(f"num decayed parameter tensors: {len(decay_params)}, with {num_decay_params:,} parameters")
        print(f"num non-decayed parameter tensors: {len(nodecay_params)}, with {num_nodecay_params:,} parameters")

        if opt_type == 'adam':
            # Create AdamW optimizer and use the fused version if it is available
            fused_available = 'fused' in inspect.signature(torch.optim.AdamW).parameters
            use_fused = fused_available and device_type == 'cuda'
            extra_args = dict(fused=True) if use_fused else dict()
            optimizer = torch.optim.AdamW(optim_groups, lr=learning_rate, betas=betas, **extra_args)
            print(f"using fused AdamW: {use_fused}")
        elif opt_type == 'gd':
            optimizer = torch.optim.SGD(optim_groups, lr=learning_rate, momentum=0)
        elif opt_type == 'gdm':
            optimizer = torch.optim.SGD(optim_groups, lr=learning_rate, momentum=betas[0], dampening=dampening)
        else:
            raise ValueError('wrong optimizer type')

        return optimizer

    def estimate_mfu(self, fwdbwd_per_iter, dt):
        """ estimate model flops utilization (MFU) in units of A100 bfloat16 peak FLOPS """
        # first estimate the number of flops we do per iteration.
        # see PaLM paper Appendix B as ref: https://arxiv.org/abs/2204.02311
        N = self.get_num_params()
        cfg = self.config
        L, H, Q, T = cfg.n_layer, cfg.n_head, cfg.n_embd//cfg.n_head, cfg.block_size
        flops_per_token = 6*N + 12*L*H*Q*T
        flops_per_fwdbwd = flops_per_token * T
        flops_per_iter = flops_per_fwdbwd * fwdbwd_per_iter
        # express our flops throughput as ratio of A100 bfloat16 peak flops
        flops_achieved = flops_per_iter * (1.0/dt) # per second
        flops_promised = 312e12 # A100 GPU bfloat16 peak flops is 312 TFLOPS
        mfu = flops_achieved / flops_promised
        return mfu

    @torch.no_grad()
    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
        """
        Take a conditioning sequence of indices idx (LongTensor of shape (b,t)) and complete
        the sequence max_new_tokens times, feeding the predictions back into the model each time.
        Most likely you'll want to make sure to be in model.eval() mode of operation for this.
        """
        for _ in range(max_new_tokens):
            # if the sequence context is growing too long we must crop it at block_size
            idx_cond = idx if idx.size(1) <= self.config.block_size else idx[:, -self.config.block_size:]
            # forward the model to get the logits for the index in the sequence
            logits, _ = self(idx_cond)
            # pluck the logits at the final step and scale by desired temperature
            logits = logits[:, -1, :] / temperature
            # optionally crop the logits to only the top k options
            if top_k is not None:
                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits[logits < v[:, [-1]]] = -float('Inf')
            # apply softmax to convert logits to (normalized) probabilities
            probs = F.softmax(logits, dim=-1)
            # sample from the distribution
            idx_next = torch.multinomial(probs, num_samples=1)
            # append sampled index to the running sequence and continue
            idx = torch.cat((idx, idx_next), dim=1)

        return idx
    
    def grad_norm(self):
        return sum([(p.grad**2).sum() for p in self.parameters()])**0.5
    
    def weight_norm(self, wd_substrings_include=['wte', 'wpe', 'mlp', 'attn', 'lm_head', 'ln']):
        squared_norm = 0
        for n, p in self.named_parameters():
            if any([s in n for s in wd_substrings_include]):
                squared_norm += (p**2).sum()
        return squared_norm**0.5
    
    def max_param_value(self):
        return max([p.abs().max() for p in self.parameters()])
    
    def min_param_value(self):
        return min([p.abs().min() for p in self.parameters()])
    
    def max_grad_value(self):
        return max([p.grad.abs().max() for p in self.parameters()])

    def compute_layerwise_norms(self):
        ln_norms, mlp_norms, att_norms, wte_norms = [], [], [], []
        for pn, p in self.named_parameters():
            if '.ln' in pn:
                ln_norms.append(torch.norm(p).item())
            elif '.mlp.' in pn:
                mlp_norms.append(torch.norm(p).item())
            elif '.attn.' in pn:
                att_norms.append(torch.norm(p).item())
            elif '.wte.' in pn:
                wte_norms.append(torch.norm(p).item())
            # print(f'{pn} ({p.shape}): {torch.norm(p).item():.2f}')
        return {
            'weights/avg_ln_norm': np.mean(ln_norms), 
            'weights/avg_mlp_norm': np.mean(mlp_norms),
            'weights/avg_att_norm': np.mean(att_norms),
            'weights/avg_wte_norm': np.mean(wte_norms),
        }
    
    def moving_average(self, model, ema_coef):
        for param1, param2 in zip(self.parameters(), model.parameters()):
            param1.data *= ema_coef
            param1.data += param2.data * (1 - ema_coef)
    

def load_state_dict(ckpt_path, device):
    checkpoint = torch.load(ckpt_path, map_location=device)
    state_dict = checkpoint['model']
    unwanted_prefix = '_orig_mod.'
    for k, v in list(state_dict.items()):
        if k.startswith(unwanted_prefix):
            state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k)
    return state_dict


large_language_models/model.py
tml-epfl/why-weight-decay
"""
This training script can be run both on a single gpu in debug mode,
and also in a larger training run with distributed data parallel (ddp).

To run on a single GPU, example:
$ python train.py --batch_size=32 --compile=False

To run with DDP on 4 gpus on 1 node, example:
$ torchrun --standalone --nproc_per_node=4 train.py

To run with DDP on 4 gpus across 2 nodes, example:
- Run on the first (master) node with example IP 123.456.123.456:
$ torchrun --nproc_per_node=8 --nnodes=2 --node_rank=0 --master_addr=123.456.123.456 --master_port=1234 train.py
- Run on the worker node:
$ torchrun --nproc_per_node=8 --nnodes=2 --node_rank=1 --master_addr=123.456.123.456 --master_port=1234 train.py
(If your cluster does not have Infiniband interconnect prepend NCCL_IB_DISABLE=1)
"""

import copy
import os
import time
import math
import pickle
from contextlib import nullcontext
from datetime import datetime

import numpy as np
import torch
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.distributed import init_process_group, destroy_process_group

from model import GPTConfig, GPT


# -----------------------------------------------------------------------------
# default config values designed to train a gpt2 (124M) on OpenWebText
# I/O
out_dir = '/home/maksym/tml_wd/models_llm'
eval_interval = 2000
log_interval = 1
eval_examples = 1000
eval_only = False # if True, script exits right after the first eval
init_from = 'scratch' # 'scratch' or 'resume' or 'gpt2*'
ckpt_path = ''
# wandb logging
wandb_log = True # disabled by default
wandb_project = 'llm'
wandb_run_name = 'owt'
# data: shakespeare: 1M of chars, owt: 9035582489 (~9B) of tokens
dataset = 'openwebtext'
gradient_accumulation_steps = 5 * 8 # used to simulate larger batch sizes
batch_size = 12 # if gradient_accumulation_steps > 1, this is the micro-batch size
block_size = 1024
# 9B tokens = 32 A100 days; 9035582489 / 32 / 24 / 60 = 196084 tokens = 1 min
# Example: 1024 context window * 256 batch size * 35000 iterations = 9,175,040,000 ≈ 9B tokens (1 epoch)
n_train_tokens = -1  
# model
random_seed = 0
n_layer = 12
n_head = 12
n_embd = 768
dropout = 0.0 # for pretraining 0 is good, for finetuning try 0.1+
bias = False # do we use bias inside LayerNorm and Linear layers?
learnable_ln = True
custom_init = False
init_scale = 0.02
weight_tying = True
force_save_ckpt = False # if True, will save a checkpoint every 2500 steps even if n_embd < 500
# adamw optimizer
opt_type = 'adam' # adam, gd, gdm
learning_rate = 6e-4 # max learning rate
max_iters = 600000 # total number of training iterations
weight_decay = 0.0
wd_substrings_include = 'wte wpe mlp attn lm_head' # apply WD only to params that contain one of these substrings
l2_reg = 0 # by default, we are not using it
beta1 = 0.9
beta2 = 0.95
dampening = 0 # for SGD+M only
grad_clip = 1.0 # clip gradients at this value, or disable if == 0.
ema_coef = 0.99
# learning rate decay settings
decay_lr = True # whether to decay the learning rate
warmup_iters = 2000 # how many steps to warm up for
lr_decay_iters = None # should be ~= max_iters per Chinchilla
min_lr = None # minimum learning rate, should be ~= learning_rate/10 per Chinchilla
# DDP settings
backend = 'nccl' # 'nccl', 'gloo', etc.
# system
device = 'cuda' # examples: 'cpu', 'cuda', 'cuda:0', 'cuda:1' etc., or try 'mps' on macbooks
dtype = 'bfloat16' if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else 'float16' # 'float32', 'bfloat16', or 'float16', the latter will auto implement a GradScaler
compile = True # use PyTorch 2.0 to compile the model to be faster
# -----------------------------------------------------------------------------
config_keys = [k for k,v in globals().items() if not k.startswith('_') and isinstance(v, (int, float, bool, str))]
exec(open('configurator.py').read()) # overrides from command line or config file
config = {k: globals()[k] for k in config_keys} # will be useful for logging
# -----------------------------------------------------------------------------
if min_lr is None:  # if min_lr is not specified in the config or as a cmd line argument
    min_lr = learning_rate / 10
if lr_decay_iters is None:
    lr_decay_iters = max_iters
wd_l2reg_param = weight_decay if weight_decay > 0.0 else l2_reg
wd_substrings_include = wd_substrings_include.split(' ')
cur_timestamp = str(datetime.now())[:-3].replace(' ', '_').replace(':', '-')
hps_str = f'{wandb_run_name}-learning_rate={learning_rate}-min_lr={min_lr:6f}-weight_decay={wd_l2reg_param}-n_embd={n_embd}-max_iters={max_iters}-init_scale={init_scale}-n_train_tokens={n_train_tokens}-seed={random_seed}'

# various inits, derived attributes, I/O setup
ddp = int(os.environ.get('RANK', -1)) != -1 # is this a ddp run?
if ddp:
    init_process_group(backend=backend)
    ddp_rank = int(os.environ['RANK'])
    ddp_local_rank = int(os.environ['LOCAL_RANK'])
    ddp_world_size = int(os.environ['WORLD_SIZE'])
    device = f'cuda:{ddp_local_rank}'
    torch.cuda.set_device(device)
    master_process = ddp_rank == 0 # this process will do logging, checkpointing etc.
    seed_offset = ddp_rank # each process gets a different seed
    # world_size number of processes will be training simultaneously, so we can scale
    # down the desired gradient accumulation iterations per process proportionally
    assert gradient_accumulation_steps % ddp_world_size == 0
    gradient_accumulation_steps //= ddp_world_size
else:
    # if not ddp, we are running on a single gpu, and one process
    master_process = True
    seed_offset = 0
    ddp_world_size = 1
tokens_per_iter = gradient_accumulation_steps * ddp_world_size * batch_size * block_size
print(f"tokens per iteration will be: {tokens_per_iter:,}")

if master_process and out_dir != '':
    os.makedirs(out_dir, exist_ok=True)
torch.manual_seed(random_seed + seed_offset)  
torch.backends.cuda.matmul.allow_tf32 = True # allow tf32 on matmul
torch.backends.cudnn.allow_tf32 = True # allow tf32 on cudnn
device_type = 'cuda' if 'cuda' in device else 'cpu' # for later use in torch.autocast
# note: float16 data type will automatically use a GradScaler
ptdtype = {'float32': torch.float32, 'bfloat16': torch.bfloat16, 'float16': torch.float16}[dtype]
ctx = nullcontext() if device_type == 'cpu' else torch.amp.autocast(device_type=device_type, dtype=ptdtype)

# poor man's data loader
data_dir = os.path.join('data', dataset)
train_data = np.memmap(os.path.join(data_dir, 'train.bin'), dtype=np.uint16, mode='r')
val_data = np.memmap(os.path.join(data_dir, 'val.bin'), dtype=np.uint16, mode='r')

def get_batch(split):
    data = train_data if split == 'train' else val_data
    if split == 'train' and n_train_tokens > 0:
        data = data[:n_train_tokens]
    ix = torch.randint(len(data) - block_size, (batch_size,))
    x = torch.stack([torch.from_numpy((data[i:i+block_size]).astype(np.int64)) for i in ix])
    y = torch.stack([torch.from_numpy((data[i+1:i+1+block_size]).astype(np.int64)) for i in ix])
    if device_type == 'cuda':
        # pin arrays x,y, which allows us to move them to GPU asynchronously (non_blocking=True)
        x, y = x.pin_memory().to(device, non_blocking=True), y.pin_memory().to(device, non_blocking=True)
    else:
        x, y = x.to(device), y.to(device)
    return x, y

# init these up here, can override if init_from='resume' (i.e. from a checkpoint)
iter_num, grad_norm, clip_times, avg_adam_elr = 0, 0, 0, 0
best_val_loss = 1e9

# attempt to derive vocab_size from the dataset
meta_path = os.path.join(data_dir, 'meta.pkl')
meta_vocab_size = None
if os.path.exists(meta_path):
    with open(meta_path, 'rb') as f:
        meta = pickle.load(f)
    meta_vocab_size = meta['vocab_size']
    print(f"found vocab_size = {meta_vocab_size} (inside {meta_path})")

# model init
model_args = dict(n_layer=n_layer, n_head=n_head, n_embd=n_embd, block_size=block_size,
                  bias=bias, learnable_ln=learnable_ln, vocab_size=None, dropout=dropout, 
                  custom_init=custom_init, init_scale=init_scale, weight_tying=weight_tying) # start with model_args from command line
if init_from == 'scratch':
    # init a new model from scratch
    print("Initializing a new model from scratch")
    # determine the vocab size we'll use for from-scratch training
    if meta_vocab_size is None:
        print("defaulting to vocab_size of GPT-2 to 50304 (50257 rounded up for efficiency)")
    model_args['vocab_size'] = meta_vocab_size if meta_vocab_size is not None else 50304
    gptconf = GPTConfig(**model_args)
    model = GPT(gptconf)
elif init_from in 'resume':
    print(f"Resuming training from {out_dir}")
    # resume training from a checkpoint
    checkpoint = torch.load(ckpt_path, map_location=device)
    checkpoint_model_args = checkpoint['model_args']
    # force these config attributes to be equal otherwise we can't even resume training
    # the rest of the attributes (e.g. dropout) can stay as desired from command line
    for k in ['n_layer', 'n_head', 'n_embd', 'block_size', 'bias', 'vocab_size']:
        model_args[k] = checkpoint_model_args[k]
    # create the model
    gptconf = GPTConfig(**model_args)
    model = GPT(gptconf)
    state_dict = checkpoint['model']
    # fix the keys of the state dictionary :(
    # honestly no idea how checkpoints sometimes get this prefix, have to debug more
    unwanted_prefix = '_orig_mod.'
    for k,v in list(state_dict.items()):
        if k.startswith(unwanted_prefix):
            state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k)
    model.load_state_dict(state_dict)
    iter_num = checkpoint['iter_num']
    max_iters += iter_num
    best_val_loss = checkpoint['best_val_loss']
elif init_from.startswith('gpt2'):
    print(f"Initializing from OpenAI GPT-2 weights: {init_from}")
    # initialize from OpenAI GPT-2 weights
    override_args = dict(dropout=dropout)
    model = GPT.from_pretrained(init_from, override_args)
    # read off the created config params, so we can store them into checkpoint correctly
    for k in ['n_layer', 'n_head', 'n_embd', 'block_size', 'bias', 'vocab_size']:
        model_args[k] = getattr(model.config, k)
# crop down the model block size if desired, using model surgery
if block_size < model.config.block_size:
    model.crop_block_size(block_size)
    model_args['block_size'] = block_size # so that the checkpoint will have the right value
model.to(device)

# initialize a GradScaler. If enabled=False scaler is a no-op
scaler = torch.cuda.amp.GradScaler(enabled=(dtype == 'float16'))

# optimizer
optimizer = model.configure_optimizers(opt_type, weight_decay, learning_rate, (beta1, beta2), device_type, wd_substrings_include, dampening)
if init_from == 'resume':
    checkpoint['optimizer']['param_groups'][0]['weight_decay'] = weight_decay
    optimizer.load_state_dict(checkpoint['optimizer'])
checkpoint = None # free up memory

# compile the model
if compile:
    print("compiling the model... (takes a ~minute)")
    unoptimized_model = model
    model = torch.compile(model) # requires PyTorch 2.0

# wrap model into DDP container
if ddp:
    model = DDP(model, device_ids=[ddp_local_rank])

model_ema = copy.deepcopy(model).eval()

# helps estimate an arbitrarily accurate loss over either split using many batches
@torch.no_grad()
def estimate_loss(model):
    metrics = {}
    model.eval()
    for split in ['train', 'val']:
        eval_iters = int(math.ceil(eval_examples / batch_size))
        losses = torch.zeros(eval_iters)
        errs = torch.zeros(eval_iters)
        for k in range(eval_iters):
            X, Y = get_batch(split)
            with ctx:
                logits, loss = model(X, Y)
            losses[k] = loss.item()
            Y_pred = logits.view(-1, logits.size(-1)).max(1)[1]
            errs[k] = (Y_pred != Y.view(-1)).float().mean()
        metrics[split+'_loss'] = losses.mean()
        metrics[split+'_err'] = errs.mean()
    model.train()
    return metrics

# learning rate decay scheduler (cosine with warmup)
def get_lr(it):
    # 1) linear warmup for warmup_iters steps
    if it < warmup_iters:
        return learning_rate * it / warmup_iters
    # 2) if it > lr_decay_iters, return min learning rate
    if it > lr_decay_iters:
        return min_lr
    # 3) in between, use cosine decay down to min learning rate
    decay_ratio = (it - warmup_iters) / (lr_decay_iters - warmup_iters)
    assert 0 <= decay_ratio <= 1
    coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio)) # coeff ranges 0..1
    return min_lr + coeff * (learning_rate - min_lr)

# logging
if wandb_log and master_process:
    import wandb
    wandb.init(project=wandb_project, name=hps_str, config=config)

# training loop
X, Y = get_batch('train') # fetch the very first batch
tstart = t0 = time.time()
local_iter_num = 0 # number of iterations in the lifetime of this process
raw_model = model.module if ddp else model # unwrap DDP container if needed
while True:
    # determine and set the learning rate for this iteration
    lr = get_lr(iter_num) if decay_lr else learning_rate
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr

    # evaluate the loss on train/val sets and write checkpoints
    if (iter_num % eval_interval == 0 or iter_num == max_iters) and master_process:
        metrics = estimate_loss(model)
        metrics_ema = estimate_loss(model_ema)

        opt_state = optimizer.state_dict()['state']
        if iter_num > 0 and opt_type == 'adam':
            num_params = sum([np.prod(param_state['exp_avg_sq'].shape) for param_state in opt_state.values()])
            avg_adam_elr = sum([torch.sum(lr / (param_state['exp_avg_sq']**0.5 + optimizer.param_groups[0]['eps'])).item() for param_state in opt_state.values()]) / num_params

        print(f"step {iter_num}: loss {metrics['train_loss']:.3f}/{metrics_ema['train_loss']:.3f}, err {metrics['train_err']:.2%}/{metrics_ema['train_err']:.2%}")
        if wandb_log:
            wandb.log({
                "iter": iter_num,
                "train/loss": metrics['train_loss'],
                "train/err": metrics['train_err'],
                "train/loss_ema": metrics_ema['train_loss'],
                "train/err_ema": metrics_ema['train_err'],
                "val/loss": metrics['val_loss'],
                "val/err": metrics['val_err'],
                "val/loss_ema": metrics_ema['val_loss'],
                "val/err_ema": metrics_ema['val_err'],
                "lr": lr,
                "avg_adam_elr": avg_adam_elr,
            })
        if metrics['val_loss'] < best_val_loss:
           best_val_loss = metrics['val_loss']
        if (iter_num % 5000 == 0 or iter_num == max_iters) and\
           (n_embd > 500 or force_save_ckpt) and\
           (init_from == 'scratch' or force_save_ckpt) and\
           out_dir != '': 
            checkpoint = {
                'model': raw_model.state_dict(),
                'optimizer': optimizer.state_dict(),
                'model_args': model_args,
                'iter_num': iter_num,
                'best_val_loss': best_val_loss,
                'config': config,
            }
            print(f"saving checkpoint to {out_dir}")
            torch.save(checkpoint, os.path.join(out_dir, f'{cur_timestamp}-{hps_str}-iter={iter_num}.pt'))
    if iter_num == 0 and eval_only:
        break

    t0 = time.time()  # start timing this iteration     

    # forward backward update, with optional gradient accumulation to simulate larger batch size
    # and using the GradScaler if data type is float16
    loss_iter, err_iter = 0, 0
    for micro_step in range(gradient_accumulation_steps):
        if ddp:
            # in DDP training we only need to sync gradients at the last micro step.
            # the official way to do this is with model.no_sync() context manager, but
            # I really dislike that this bloats the code and forces us to repeat code
            # looking at the source of that context manager, it just toggles this variable
            model.require_backward_grad_sync = (micro_step == gradient_accumulation_steps - 1)
        with ctx:
            logits, loss = model(X, Y)
            loss = loss / gradient_accumulation_steps # scale the loss to account for gradient accumulation
            Y_pred = logits.view(-1, logits.size(-1)).max(1)[1]
            err = (Y_pred != Y.view(-1)).float().mean() / gradient_accumulation_steps
            loss_iter += loss.item()
            err_iter += err.item()
        # immediately async prefetch next batch while model is doing the forward pass on the GPU
        X, Y = get_batch('train')
        # backward pass, with gradient scaling if training in fp16
        scaler.scale(loss).backward()
    
    if np.isnan(loss_iter):
        print("loss is NaN, aborting training")
        break

    if l2_reg > 0:
        assert weight_decay == 0.0  # make sure we are using only l2_reg and not WD
        reg = l2_reg * model.weight_norm(wd_substrings_include)**2
        scaler.scale(reg).backward()
    
    grad_max_value = raw_model.max_grad_value()
    grad_layerwise_norms = {'grads_layerwise/' + '.'.join(pn.split('.')[2:-1]): n.grad.norm().item() for pn, n in model.named_parameters()}

    # clip the gradient
    if grad_clip > 0.0:
        scaler.unscale_(optimizer)
        grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
        if grad_norm.item() > grad_clip:
            clip_times += 1
    else:
        grad_norm = raw_model.grad_norm()
    # step the optimizer and scaler if training in fp16
    scaler.step(optimizer)
    scaler.update()
    # flush the gradients as soon as we can, no need for this memory anymore
    optimizer.zero_grad(set_to_none=True)

    model_ema.moving_average(model, ema_coef)

    if wandb_log and iter_num % 20 == 0:  # log every K steps so that the legends are properly displayed on hover
        layerwise_norms = model.compute_layerwise_norms()
        wandb.log({
            "iter": iter_num,
            "train_every_iter/loss": loss_iter,
            "train_every_iter/err": err_iter,
            "grads/scaler_scale": scaler.get_scale(),  
            "grads/grad_norm": grad_norm,
            "grads/grad_max_value": grad_max_value,
            "grads/clip_times": clip_times,
            "weights/weight_norm": raw_model.weight_norm().item(),
            "weights/max_param_value": raw_model.max_param_value().item(),
            **layerwise_norms,
            **grad_layerwise_norms,
        })
    if iter_num % log_interval == 0 and master_process:
        print(f"iter {iter_num}: loss {loss_iter:.3f}, err {err_iter:.2%}, iter time {time.time()-t0:.3f}, total time {(time.time()-tstart)/60:.2f}m")
    iter_num += 1
    local_iter_num += 1

    # termination conditions
    if iter_num > max_iters:
        break

if ddp:
    destroy_process_group()


large_language_models/train.py
tml-epfl/why-weight-decay
from datetime import datetime
import argparse
from data import get_dataset

dout_dir = (
    "./out/"
    + datetime.now().strftime("%Y-%m-%d")
    + "/run_"
    + datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
)

default_args = dict(
    epochs=100,
    lr=0.03,
    random_seed=42,
    batch_size = 128,
    save_every = 100,
    out_dir=dout_dir,
    date_time_dir = None, 
    exp_name="resnet_test",
    wd = 0.0,
    use_cuda=True,
    det_run=True,
    scale_inv = False,
    loglevel_info=False,
    model = 'resnet18',
    model_width = 64, #default for VGG
    half_prec=False,
    no_data_augm = True,
    batch_norm = False,
    dataset = 'cifar10',
    momentum = 0.9,
    first_decay = 0.5,
    lr_gamma_decay = 1.0, 
    lr_flow=1e-4,
    flow_steps=100,
    flow_every = 50,
    wd_flow=0.0,
    wd_gamma_decay = 1.0,
    ckpt = None, 
    config_file = None,
    radius = None
)

def configuration(args=None):
    parser = argparse.ArgumentParser(description="Exp Resnet")

    #--------------------------------------------------------------------------------
    # TRAINING OPTIONS 
    #--------------------------------------------------------------------------------
    tgroup = parser.add_argument_group("Training options")
    tgroup.add_argument(
        "--epochs",
        type=int,
        metavar="N",
        help="Number of training iters. " + "Default: %(default)s.",
    )
    tgroup.add_argument(
        "--flow_steps",
        type=int,
        metavar="N",
        help="Number of gradient flow iters. " + "Default: %(default)s.",
    )
    tgroup.add_argument(
        "--flow_every",
        type=int,
        metavar="N",
        help="Lr is decayed at multiples of this number. " + "Default: %(default)s.",
    )
    tgroup.add_argument(
        "--batch_size",
        type=int,
        metavar="N",
        help="Batch size. " + "Default: %(default)s.",
    )
    tgroup.add_argument(
        "--lr",
        type=float,
        help="Learning rate of optimizer. Default: " + "%(default)s.",
    )
    tgroup.add_argument(
        "--first_decay",
        type=float,
        help="Fraction of epochs when lr is first decayed: " + "%(default)s.",
    )
    tgroup.add_argument(
        "--lr_gamma_decay",
        type=float,
        help="Decaying factor of the learning rate: " + "%(default)s.",
    )
    tgroup.add_argument(
        "--lr_flow",
        type=float,
        help="Decaying the learning rate to this exact value: " +
        "%(default)s.",
    )
    tgroup.add_argument(
        "--wd_flow",
        type=float,
        help="Decaying the learning rate to this exact value: " +
        "%(default)s.",
    )
    tgroup.add_argument(
        "--wd_gamma_decay",
        type=float,
        help="Decaying factor of the weight decay: " + "%(default)s.",
    )
    tgroup.add_argument(
        "--wd",
        type=float,
        help="Weight decay of optimizer. Default: " + "%(default)s.",
    )

    tgroup.add_argument(
        "--radius",
        type=float,
        help="Radius for PGD. Default: " + "%(default)s.",
    )
    tgroup.add_argument('--momentum', type=float, help="Momentum of optimizer. Default: " + "%(default)s.",)

    #--------------------------------------------------------------------------------
    # MODEL OPTIONS 
    #--------------------------------------------------------------------------------
    modelgroup = parser.add_argument_group("Model options")
    modelgroup.add_argument('--model', choices=['vit_basic', 'vit_exp', 'vgg16', 'resnet18', 'resnet18_plain', 'resnet18_gn', 'resnet_tiny', 'resnet_tiny_gn', 'resnet34', 'resnet34_plain', 'resnet34preact', 'resnet34_gn', 'wrn28', 'lenet', 'cnn', 'fc', 'linear'], type=str)
    modelgroup.add_argument('--model_width', type=int, help='model width (# conv filters on the first layer for ResNets)')
    modelgroup.add_argument('--half_prec', action='store_true', help='if enabled, runs everything as half precision [not recommended]')
    modelgroup.add_argument('--batch_norm', action = 'store_true', help = 'if enabled, adds batch-norm layers')
    modelgroup.add_argument('--scale_inv', action = 'store_true', help = 'if enabled, makes the model scale invariant, only for resnets')

    #--------------------------------------------------------------------------------
    # MISCELLANEOUS OPTIONS 
    #--------------------------------------------------------------------------------
    mgroup = parser.add_argument_group("Miscellaneous options")
    mgroup.add_argument(
        "--random_seed",
        type=int,
        metavar="N",
        help="Random seed. Default: %(default)s.",
    )
    mgroup.add_argument(
        "--save_every",
        type=int,
        metavar="N",
        help="Numer of iterations for saving model. Default: %(default)s.",
    )
    mgroup.add_argument(
        "--out_dir", type=str, help="Where to store the outputs of this simulation."
    )
    mgroup.add_argument("--exp_name", type=str, help="Name of the experiment.")
    
    mgroup.add_argument("--date_time_dir", type=str, help="Name of the date time dir.")

    
    #--------------------------------------------------------------------------------
    # EXPERIMENTS OPTIONS 
    #--------------------------------------------------------------------------------
    expgroup = parser.add_argument_group("Experiments options")
    expgroup.add_argument(
        "--use_cuda",
        action="store_true",
        help="Use cuda",
    )
    expgroup.add_argument(
        "--det_run",
        action="store_true",
        help="Ensure deterministic behaviour",
    )
    expgroup.add_argument(
        "--loglevel_info",
        action="store_true",
        help="Info in the logger",
    )
    expgroup.add_argument(
        "--ckpt",
        type=str,
        help="Path for model checkpoint: " + "%(default)s.",
    )
    expgroup.add_argument(
        "--config_file",
        type=str,
        help="Path to specify config file: " + "%(default)s.",
    )
    expgroup.add_argument('--no_data_augm', action='store_true')

    #--------------------------------------------------------------------------------
    # DATA OPTIONS 
    #--------------------------------------------------------------------------------
    datagroup = parser.add_argument_group("Dataset options")
    datagroup.add_argument(
        "--input_dim",
        type=int,
        help="Random seed. Default: %(default)s.",
    )
    datagroup.add_argument('--dataset', choices=get_dataset.datasets_dict.keys(), type=str)

    parser.set_defaults(**default_args)

    if args is None:
        return parser.parse_args()
    else:
        return parser.parse_args(args=args)

overparameterized_nets/configs/config.py
tml-epfl/why-weight-decay
from typing import List
import os 
import torch as ch
import torchvision
from ffcv.writer import DatasetWriter
from ffcv.fields import IntField, RGBImageField
import numpy as np 
from ffcv.fields.decoders import IntDecoder, SimpleRGBImageDecoder
from ffcv.loader import Loader, OrderOption
from ffcv.pipeline.operation import Operation
from ffcv.transforms import RandomHorizontalFlip, Cutout, \
    RandomTranslate, Convert, ToDevice, ToTensor, ToTorchImage
from ffcv.transforms.common import Squeeze
from ffcv.writer import DatasetWriter
from torchvision import datasets
# Note that statistics are wrt to uin8 range, [0,255].
CIFAR_MEAN = [125.307, 122.961, 113.8575]
CIFAR_STD = [51.5865, 50.847, 51.255]

def create_dataloaders(dataset, no_data_augm , batch_size, device):
    loaders = {}

    # Create loaders
    if dataset == 'cifar10': 
        path = f'data/CIFAR10/cifar10'
        if not os.path.isfile(path+f'_train.beton'): 
            data = {
                'train': torchvision.datasets.CIFAR10('data/CIFAR10', train=True, download=True),
                'test': torchvision.datasets.CIFAR10('data/CIFAR10', train=False, download=True)}
            for (name, ds) in data.items():
                writer = DatasetWriter(f'data/CIFAR10/cifar10_{name}.beton', {
                    'image': RGBImageField(),
                    'label': IntField()
                })
                writer.from_indexed_dataset(ds)

    elif dataset == 'cifar10_5k':
        path = f'data/CIFAR10/cifar10_5k'
        if not os.path.isfile(path+f'_train.beton'): 
            data = {
                'train': torchvision.datasets.CIFAR10('data/CIFAR10', train=True, download=True),
                'test': torchvision.datasets.CIFAR10('data/CIFAR10', train=False, download=True)}
            n_points = 500 #n_points per class 
            new_train_data = []
            new_train_labels = []
            for i in range(10):
                indices = np.where(np.array(data['train'].targets) == i)[0][:n_points]
                new_train_data.append(data['train'].data[indices])
                new_train_labels += [i]*n_points
            new_train_data = np.concatenate(new_train_data, axis=0)
            data['train'].data = new_train_data
            data['train'].targets = new_train_labels
            for (name, ds) in data.items():
                writer = DatasetWriter(f'data/CIFAR10/cifar10_5k_{name}.beton', {
                    'image': RGBImageField(),
                    'label': IntField()
                })
                writer.from_indexed_dataset(ds)

    elif dataset == 'cifar100':
        path = f'data/CIFAR100/cifar100'
        if not os.path.isfile(path+f'_train.beton'): 
            data = {
                'train': torchvision.datasets.CIFAR100('data/CIFAR100', train=True, download=True),
                'test': torchvision.datasets.CIFAR100('data/CIFAR100', train=False, download=True)}
            for (name, ds) in data.items():
                writer = DatasetWriter(f'data/CIFAR100/cifar100_{name}.beton', {
                    'image': RGBImageField(),
                    'label': IntField()
                })
                writer.from_indexed_dataset(ds)



    for name in ['train', 'test']:
        label_pipeline: List[Operation] = [IntDecoder(), ToTensor(), ToDevice(device), Squeeze()]
        image_pipeline: List[Operation] = [SimpleRGBImageDecoder()]

        # Add image transforms and normalization
        if name == 'train' and not no_data_augm:
            image_pipeline.extend([
                RandomHorizontalFlip(),
                RandomTranslate(padding=2),
                Cutout(8, tuple(map(int, CIFAR_MEAN))), # Note Cutout is done before normalization.
            ])
        image_pipeline.extend([
            ToTensor(),
            ToDevice(device, non_blocking=True),
            ToTorchImage(),
            Convert(ch.float16),
            torchvision.transforms.Normalize(CIFAR_MEAN, CIFAR_STD),
        ])

        loaders[name] = Loader(path+f'_{name}.beton',
                                batch_size=batch_size,
                                num_workers=20,
                                order=OrderOption.RANDOM,
                                drop_last=(name == 'train'),
                                pipelines={'image': image_pipeline,
                                           'label': label_pipeline})
    return loaders 


shapes_dict = {'mnist': (60000, 1, 28, 28),
               'mnist_binary': (13007, 1, 28, 28),
               'svhn': (73257, 3, 32, 32),
               'cifar10': (50000, 3, 32, 32),
               'cifar10_horse_car': (10000, 3, 32, 32),
               'cifar10_binary': (10000, 3, 32, 32),
               'cifar10_10000': (10000, 3, 32, 32),
               'cifar10_dog_cat': (10000, 3, 32, 32),
               'cifar100': (50000, 3, 32, 32),
               'uniform_noise': (1000, 1, 28, 28),
               'gaussians_binary': (1000, 1, 1, 100),
               }

datasets_dict = {'mnist': datasets.MNIST,
                 'mnist_binary': datasets.MNIST,
                 'svhn': datasets.SVHN,
                 'cifar10': datasets.CIFAR10,
                 'cifar10_horse_car': datasets.CIFAR10,
                 'cifar10_binary': datasets.CIFAR10,
                 'cifar10_10000': datasets.CIFAR10,
                 'cifar10_dog_cat': datasets.CIFAR10,
                 'cifar100': datasets.CIFAR100,
                 }

classes_dict = {'cifar10': {0: 'airplane',
                            1: 'automobile',
                            2: 'bird',
                            3: 'cat',
                            4: 'deer',
                            5: 'dog',
                            6: 'frog',
                            7: 'horse',
                            8: 'ship',
                            9: 'truck',
                            }
                }


                
overparameterized_nets/data/get_dataset.py
tml-epfl/why-weight-decay
import os
import random
from datetime import datetime
import numpy as np
import torch
import yaml
from prettyprinter import pprint
from aim import Run


def set_exp(config):
    """This utils sets up all the experimental artifacts such as random seeds,
       log folders for results, tensorboard and aim loggers, saves a yaml file
       with all the hyperparams.

    Args:
        config: configuratore for the experiment

    Returns:
        device: torch devide
        aim_run: aim writer
    """

    experiment_name = config.exp_name
    # Create Aim writer
    aim_writer = Run(
        repo='/tmldata1/fdangelo/understanding-weight-decay/aim_paper_2', experiment=experiment_name)
    # Create experiments folders and artifacts
    if config.date_time_dir is None:
        date_time_log_dir = datetime.now().strftime("%m-%d-%Y")
    else:
        date_time_log_dir = config.date_time_dir
    subdirname = aim_writer.hash
    log_dir = os.path.join(
        os.getcwd(),
        "out",
        experiment_name,
        date_time_log_dir,
        subdirname,
    )

    os.makedirs(log_dir)
    os.makedirs(os.path.join(log_dir, "models"))
    config.out_dir = log_dir
    print(f"Experiment folder created at {log_dir}")

    # Save parameters file
    configu = dict(config.__dict__)
    configu['hash'] = aim_writer.hash
    # Add Hparams to tensorbaord
    sorted_dict = {key: value for key, value in sorted(configu.items())}
    aim_writer['hparams'] = sorted_dict
    pprint(sorted_dict)

    dictionary_parameters = vars(config)

    # Deterministic computation.
    torch.manual_seed(config.random_seed)
    torch.cuda.manual_seed_all(config.random_seed)
    np.random.seed(config.random_seed)
    random.seed(config.random_seed)

    # Ensure that runs are reproducible. Note, this slows down training!
    # https://pytorch.org/docs/stable/notes/randomness.html
    if config.det_run:
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False

    # Select torch device.
    assert (hasattr(config, 'no_cuda') or hasattr(config, 'use_cuda'))
    assert (not hasattr(config, 'no_cuda') or not hasattr(config, 'use_cuda'))

    if hasattr(config, 'no_cuda'):
        use_cuda = not config.no_cuda and torch.cuda.is_available()
    else:
        use_cuda = config.use_cuda and torch.cuda.is_available()
    device = torch.device("cuda" if use_cuda else "cpu")

    with open(log_dir + "/" + "parameters.yml", "w") as yaml_file:
        yaml.dump(dictionary_parameters, stream=yaml_file,
                  default_flow_style=False)

    return device, aim_writer

overparameterized_nets/exp_utils/setup_exp.py
tml-epfl/why-weight-decay
import torch.optim.lr_scheduler as lr_scheduler


class CustomMultiStepLR(lr_scheduler._LRScheduler):
    def __init__(self, optimizer, milestones, gamma_lr=0.1, gamma_wd=0.1, last_epoch=-1):
        self.milestones = milestones
        self.gamma_lr = gamma_lr
        self.gamma_wd = gamma_wd
        self.base_wds = [group['weight_decay'] for group in optimizer.param_groups]
        super(CustomMultiStepLR, self).__init__(optimizer, last_epoch)

    def get_lr(self):
        lr_factor = 1.0
        wd_factor = 1.0
        for milestone in self.milestones:
            if self.last_epoch >= milestone:
                lr_factor *= self.gamma_lr
                wd_factor *= self.gamma_wd
        return [base_lr * lr_factor for base_lr in self.base_lrs], [base_wd * wd_factor for base_wd in self.base_wds]
    
    def step(self, epoch=None):
        if epoch is None:
            epoch = self.last_epoch + 1
        self.last_epoch = epoch
        lr, wd = self.get_lr()
        for i, param_group in enumerate(self.optimizer.param_groups):
            param_group['lr'] = lr[i]
            param_group['weight_decay'] = wd[i]
overparameterized_nets/exp_utils/utils.py
tml-epfl/why-weight-decay
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import math
import os


class Flatten(nn.Module):
    def forward(self, x):
        return x.view(x.size(0), -1)


class Normalize(nn.Module):
    def __init__(self, mu, std):
        super(Normalize, self).__init__()
        self.mu, self.std = mu, std

    def forward(self, x):
        return (x - self.mu) / self.std


class CustomReLU(nn.Module):
    def __init__(self):
        super(CustomReLU, self).__init__()
        self.collect_preact = True
        self.avg_preacts = []

    def forward(self, preact):
        if self.collect_preact:
            self.avg_preacts.append(preact.abs().mean().item())
        act = F.relu(preact)
        return act


class ModuleWithStats(nn.Module):
    def __init__(self):
        super(ModuleWithStats, self).__init__()

    def forward(self, x):
        for layer in self._model:
            if type(layer) == CustomReLU:
                layer.avg_preacts = []

        out = self._model(x)

        avg_preacts_all = [
            layer.avg_preacts for layer in self._model if type(layer) == CustomReLU
        ]
        self.avg_preact = np.mean(avg_preacts_all)
        return out


class Linear(ModuleWithStats):
    def __init__(self, n_cls, shape_in):
        n_cls = 1 if n_cls == 2 else n_cls
        super(Linear, self).__init__()
        d = int(np.prod(shape_in[1:]))
        self._model = nn.Sequential(Flatten(), nn.Linear(d, n_cls, bias=False))

    def forward(self, x):
        logits = self._model(x)
        return torch.cat([torch.zeros(logits.shape).cuda(), logits], dim=1)


class LinearTwoOutputs(ModuleWithStats):
    def __init__(self, n_cls, shape_in):
        super(LinearTwoOutputs, self).__init__()
        d = int(np.prod(shape_in[1:]))
        self._model = nn.Sequential(Flatten(), nn.Linear(d, n_cls, bias=False))


class IdentityLayer(nn.Module):
    def forward(self, inputs):
        return inputs


class PreActBlock(nn.Module):
    """Pre-activation version of the BasicBlock."""

    expansion = 1

    def __init__(
        self,
        in_planes,
        planes,
        bn,
        learnable_bn,
        stride=1,
        activation="relu",
        droprate=0.0,
        gn_groups=32,
    ):
        super(PreActBlock, self).__init__()
        self.collect_preact = True
        self.activation = activation
        self.droprate = droprate
        self.avg_preacts = []
        self.bn1 = (
            nn.BatchNorm2d(in_planes, affine=learnable_bn)
            if bn
            else nn.GroupNorm(gn_groups, in_planes)
        )
        self.conv1 = nn.Conv2d(
            in_planes,
            planes,
            kernel_size=3,
            stride=stride,
            padding=1,
            bias=not learnable_bn,
        )
        self.bn2 = (
            nn.BatchNorm2d(planes, affine=learnable_bn)
            if bn
            else nn.GroupNorm(gn_groups, planes)
        )
        self.conv2 = nn.Conv2d(
            planes, planes, kernel_size=3, stride=1, padding=1, bias=not learnable_bn
        )

        if stride != 1 or in_planes != self.expansion * planes:
            self.shortcut = nn.Sequential(
                # We added this line to ensure scale invariance
                # nn.BatchNorm2d(in_planes, affine=learnable_bn),
                nn.Conv2d(
                    in_planes,
                    self.expansion * planes,
                    kernel_size=1,
                    stride=stride,
                    bias=not learnable_bn,
                )
            )

    def act_function(self, preact):
        if self.activation == "relu":
            act = F.relu(preact)
            # print((act == 0).float().mean().item(), (act.norm() / act.shape[0]).item(), (act.norm() / np.prod(act.shape)).item())
        else:
            assert self.activation[:8] == "softplus"
            beta = int(self.activation.split("softplus")[1])
            act = F.softplus(preact, beta=beta)
        return act

    def forward(self, x):
        out = self.act_function(self.bn1(x))
        shortcut = (
            self.shortcut(out) if hasattr(self, "shortcut") else x
        )  # Important: using out instead of x
        out = self.conv1(out)
        out = self.act_function(self.bn2(out))
        if self.droprate > 0:
            out = F.dropout(out, p=self.droprate, training=self.training)
        out = self.conv2(out)
        out += shortcut
        return out


class BasicBlock(nn.Module):
    def __init__(self, in_planes, out_planes, stride, droprate=0.0):
        super(BasicBlock, self).__init__()
        self.bn1 = nn.BatchNorm2d(in_planes)
        self.relu1 = nn.ReLU(inplace=True)
        self.conv1 = nn.Conv2d(
            in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False
        )
        self.bn2 = nn.BatchNorm2d(out_planes)
        self.relu2 = nn.ReLU(inplace=True)
        self.conv2 = nn.Conv2d(
            out_planes, out_planes, kernel_size=3, stride=1, padding=1, bias=False
        )
        self.droprate = droprate
        self.equalInOut = in_planes == out_planes
        self.convShortcut = (
            (not self.equalInOut)
            and nn.Conv2d(
                in_planes,
                out_planes,
                kernel_size=1,
                stride=stride,
                padding=0,
                bias=False,
            )
            or None
        )

    def forward(self, x):
        if not self.equalInOut:
            x = self.relu1(self.bn1(x))
        else:
            out = self.relu1(self.bn1(x))
        out = self.relu2(self.bn2(self.conv1(out if self.equalInOut else x)))
        if self.droprate > 0:
            out = F.dropout(out, p=self.droprate, training=self.training)
        out = self.conv2(out)
        return torch.add(x if self.equalInOut else self.convShortcut(x), out)


class BasicBlockResNet34(nn.Module):
    expansion = 1

    def __init__(self, in_planes, planes, stride=1):
        super(BasicBlockResNet34, self).__init__()
        self.conv1 = nn.Conv2d(
            in_planes, planes, kernel_size=3, stride=stride, padding=1, bias=False
        )
        self.bn1 = nn.BatchNorm2d(planes)
        self.conv2 = nn.Conv2d(
            planes, planes, kernel_size=3, stride=1, padding=1, bias=False
        )
        self.bn2 = nn.BatchNorm2d(planes)

        self.shortcut = nn.Sequential()
        if stride != 1 or in_planes != self.expansion * planes:
            self.shortcut = nn.Sequential(
                nn.Conv2d(
                    in_planes,
                    self.expansion * planes,
                    kernel_size=1,
                    stride=stride,
                    bias=False,
                ),
                nn.BatchNorm2d(self.expansion * planes),
            )

    def forward(self, x):
        out = F.relu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))
        out += self.shortcut(x)
        out = F.relu(out)
        return out


class NetworkBlock(nn.Module):
    def __init__(self, nb_layers, in_planes, out_planes, block, stride, droprate=0.0):
        super(NetworkBlock, self).__init__()
        self.layer = self._make_layer(
            block, in_planes, out_planes, nb_layers, stride, droprate
        )

    def _make_layer(self, block, in_planes, out_planes, nb_layers, stride, droprate):
        layers = []
        for i in range(int(nb_layers)):
            layers.append(
                block(
                    i == 0 and in_planes or out_planes,
                    out_planes,
                    i == 0 and stride or 1,
                    droprate,
                )
            )
        return nn.Sequential(*layers)

    def forward(self, x):
        return self.layer(x)


class ResNet(nn.Module):
    def __init__(self, block, num_blocks, num_classes=10, model_width=64, droprate=0.0):
        super(ResNet, self).__init__()
        self.in_planes = model_width
        self.half_prec = False
        # self.mu = torch.tensor((0.4914, 0.4822, 0.4465)).view(1, 3, 1, 1).cuda()
        # self.std = torch.tensor((0.2471, 0.2435, 0.2616)).view(1, 3, 1, 1).cuda()
        self.mu = torch.tensor((0.0, 0.0, 0.0)).view(1, 3, 1, 1).cuda()
        self.std = torch.tensor((1.0, 1.0, 1.0)).view(1, 3, 1, 1).cuda()
        # if self.half_prec:
        #     self.mu, self.std = self.mu.half(), self.std.half()

        self.normalize = Normalize(self.mu, self.std)
        self.conv1 = nn.Conv2d(
            3, model_width, kernel_size=3, stride=1, padding=1, bias=False
        )
        self.bn1 = nn.BatchNorm2d(model_width)
        self.layer1 = self._make_layer(block, model_width, num_blocks[0], stride=1)
        self.layer2 = self._make_layer(block, 2 * model_width, num_blocks[1], stride=2)
        self.layer3 = self._make_layer(block, 4 * model_width, num_blocks[2], stride=2)
        self.layer4 = self._make_layer(block, 8 * model_width, num_blocks[3], stride=2)
        self.linear = nn.Linear(8 * model_width * block.expansion, num_classes)

    def _make_layer(self, block, planes, num_blocks, stride):
        strides = [stride] + [1] * (num_blocks - 1)
        layers = []
        for stride in strides:
            layers.append(block(self.in_planes, planes, stride))
            self.in_planes = planes * block.expansion
        return nn.Sequential(*layers)

    def forward(self, x, return_features=False, return_block=5):
        assert return_block in [1, 2, 3, 4, 5], "wrong return_block"
        # out = self.normalize(x)
        out = F.relu(self.bn1(self.conv1(x)))
        out = self.layer1(out)
        if return_features and return_block == 1:
            return out
        out = self.layer2(out)
        if return_features and return_block == 2:
            return out
        out = self.layer3(out)
        if return_features and return_block == 3:
            return out
        out = self.layer4(out)
        if return_features and return_block == 4:
            return out
        out = F.avg_pool2d(out, 4)
        out = out.view(out.size(0), -1)
        if return_features and return_block == 5:
            return out
        out = self.linear(out)
        return out


class PreActResNet(nn.Module):
    def __init__(
        self,
        block,
        num_blocks,
        n_cls,
        model_width=64,
        cuda=True,
        half_prec=False,
        activation="relu",
        droprate=0.0,
        bn_flag=True,
        learnable_bn=True,
        freeze_last_layer=False,
    ):
        super(PreActResNet, self).__init__()
        self.half_prec = half_prec
        self.bn_flag = bn_flag
        # in particular, 32 for model_width=64 as in the original GroupNorm paper
        self.gn_groups = model_width // 2
        self.learnable_bn = learnable_bn  # doesn't matter if self.bn=False
        self.freeze_last_layer = freeze_last_layer
        self.in_planes = model_width
        self.avg_preact = None
        self.activation = activation
        self.n_cls = n_cls
        # self.mu = torch.tensor((0.4914, 0.4822, 0.4465)).view(1, 3, 1, 1)
        # self.std = torch.tensor((0.2471, 0.2435, 0.2616)).view(1, 3, 1, 1)
        self.mu = torch.tensor((0.0, 0.0, 0.0)).view(1, 3, 1, 1)
        self.std = torch.tensor((1.0, 1.0, 1.0)).view(1, 3, 1, 1)

        if cuda:
            self.mu, self.std = self.mu.cuda(), self.std.cuda()
        # if half_prec:
        #     self.mu, self.std = self.mu.half(), self.std.half()

        self.normalize = Normalize(self.mu, self.std)
        self.bn1 = (
            nn.BatchNorm2d(3, affine=self.learnable_bn)
            if self.bn_flag
            else nn.GroupNorm(self.gn_groups, 3)
        )
        self.conv1 = nn.Conv2d(
            3,
            model_width,
            kernel_size=3,
            stride=1,
            padding=1,
            bias=not self.learnable_bn,
        )
        self.layer1 = self._make_layer(block, model_width, num_blocks[0], 1, droprate)
        self.layer2 = self._make_layer(
            block, 2 * model_width, num_blocks[1], 2, droprate
        )
        self.layer3 = self._make_layer(
            block, 4 * model_width, num_blocks[2], 2, droprate
        )
        final_layer_factor = 8
        self.layer4 = self._make_layer(
            block, final_layer_factor * model_width, num_blocks[3], 2, droprate
        )
        self.bn = (
            nn.BatchNorm2d(
                final_layer_factor * model_width * block.expansion,
                affine=self.learnable_bn,
            )
            if self.bn_flag
            else nn.GroupNorm(
                self.gn_groups, final_layer_factor * model_width * block.expansion
            )
        )
        self.linear = nn.Linear(
            final_layer_factor * model_width * block.expansion,
            1 if n_cls == 2 else n_cls,
        )

        if self.freeze_last_layer:
            # We freeze the last linear layer to ensure scale invariance
            self.linear.weight.requires_grad = False
            # We freeze the last linear layer to ensure scale invariance
            self.linear.bias.requires_grad = False

    def _make_layer(self, block, planes, num_blocks, stride, droprate):
        strides = [stride] + [1] * (num_blocks - 1)
        layers = []
        for stride in strides:
            layers.append(
                block(
                    self.in_planes,
                    planes,
                    self.bn_flag,
                    self.learnable_bn,
                    stride,
                    self.activation,
                    droprate,
                    self.gn_groups,
                )
            )
            # layers.append(block(self.in_planes, planes, stride))
            self.in_planes = planes * block.expansion
        return nn.Sequential(*layers)

    def forward(self, x, return_features=False, return_block=5):
        assert return_block in [1, 2, 3, 4, 5], "wrong return_block"
        for layer in [*self.layer1, *self.layer2, *self.layer3, *self.layer4]:
            layer.avg_preacts = []

        # x = x / ((x**2).sum([1, 2, 3], keepdims=True)**0.5 + 1e-6)  # numerical stability is needed for RLAT
        out = self.normalize(x)
        # We added this following Arora's paper about exp learning rate
        out = self.bn1(out)
        out = self.conv1(out)
        out = self.layer1(out)
        out = self.layer2(out)
        out_3 = self.layer3(out)
        out_4 = self.layer4(out_3)
        out_4 = F.relu(self.bn(out_4))
        out_5 = F.avg_pool2d(out_4, 4)
        out_5 = out_5.view(out_5.size(0), -1)

        out = self.linear(out_5)
        if out.shape[1] == 1:
            out = torch.cat([torch.zeros_like(out), out], dim=1)
        if return_features:
            return out, out_3, out_4, out_5
        else:
            return out

    def moving_average(self, model, ema_coef):
        for param1, param2 in zip(self.parameters(), model.parameters()):
            param1.data *= ema_coef
            param1.data += param2.data * (1 - ema_coef)


class WideResNet(nn.Module):
    """Based on code from https://github.com/yaodongyu/TRADES"""

    def __init__(
        self, depth=28, num_classes=10, widen_factor=10, droprate=0.0, bias_last=True
    ):
        super(WideResNet, self).__init__()
        self.half_prec = False
        nChannels = [16, 16 * widen_factor, 32 * widen_factor, 64 * widen_factor]
        assert (depth - 4) % 6 == 0
        n = (depth - 4) / 6
        block = BasicBlock
        # 1st conv before any network block
        self.conv1 = nn.Conv2d(
            3, nChannels[0], kernel_size=3, stride=1, padding=1, bias=False
        )
        # 1st block
        self.block1 = NetworkBlock(n, nChannels[0], nChannels[1], block, 1, droprate)
        # 2nd block
        self.block2 = NetworkBlock(n, nChannels[1], nChannels[2], block, 2, droprate)
        # 3rd block
        self.block3 = NetworkBlock(n, nChannels[2], nChannels[3], block, 2, droprate)
        # global average pooling and classifier
        self.bn1 = nn.BatchNorm2d(nChannels[3])
        self.relu = nn.ReLU(inplace=True)
        self.fc = nn.Linear(nChannels[3], num_classes, bias=bias_last)
        self.nChannels = nChannels[3]

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2.0 / n))
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()
            elif isinstance(m, nn.Linear) and not m.bias is None:
                m.bias.data.zero_()

    def forward(self, x):
        out = self.conv1(x)
        out = self.block1(out)
        out = self.block2(out)
        out = self.block3(out)
        out = self.relu(self.bn1(out))
        out = F.avg_pool2d(out, 8)
        out = out.view(-1, self.nChannels)
        return self.fc(out)


class VGG(nn.Module):
    """
    VGG model. Source: https://github.com/chengyangfu/pytorch-vgg-cifar10/blob/master/vgg.py
    (in turn modified from https://github.com/pytorch/vision.git)
    """

    def __init__(self, n_cls, half_prec, cfg, batch_norm, bias):
        super(VGG, self).__init__()
        self.half_prec = half_prec
        self.mu = torch.tensor((0.485, 0.456, 0.406)).view(1, 3, 1, 1).cuda()
        self.std = torch.tensor((0.229, 0.224, 0.225)).view(1, 3, 1, 1).cuda()
        self.normalize = Normalize(self.mu, self.std)
        self.features = self.make_layers(cfg, batch_norm=batch_norm, bias=bias)
        n_out = cfg[-2]  # equal to 8*model_width
        self.classifier = nn.Sequential(
            # nn.Dropout(),
            nn.Linear(n_out, n_out, bias=bias),
            nn.ReLU(True),
            # nn.Dropout(),
            nn.Linear(n_out, n_out, bias=bias),
            nn.ReLU(True),
            nn.Linear(n_out, n_cls, bias=bias),
        )
        # Initialize weights
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2.0 / n))
                print(bias)
                if bias:
                    m.bias.data.zero_()

    def make_layers(self, cfg, batch_norm=False, bias=False):
        layers = []
        in_channels = 3
        for v in cfg:
            if v == "M":
                layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
            else:
                conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1, bias=bias)
                if batch_norm:
                    layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)]
                else:
                    layers += [conv2d, nn.ReLU(inplace=True)]
                in_channels = v
        return nn.Sequential(*layers)

    def forward(self, x):
        x = self.normalize(x)
        x = self.features(x)
        x = x.view(x.size(0), -1)
        x = self.classifier(x)
        return x

    def moving_average(self, model, ema_coef):
        for param1, param2 in zip(self.parameters(), model.parameters()):
            param1.data *= ema_coef
            param1.data += param2.data * (1 - ema_coef)


def VGG16(n_cls, model_width, half_prec, batch_norm, bias):
    """VGG 16-layer model (configuration "D")"""
    w1, w2, w3, w4, w5 = (
        model_width,
        2 * model_width,
        4 * model_width,
        8 * model_width,
        8 * model_width,
    )
    # cfg = [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M']
    cfg = [w1, w1, "M", w2, w2, "M", w3, w3, w3, "M", w4, w4, w4, "M", w5, w5, w5, "M"]
    return VGG(n_cls, half_prec, cfg, batch_norm, bias)


def TinyResNet(
    n_cls, model_width=64, cuda=True, half_prec=False, activation="relu", droprate=0.0
):
    bn_flag = True
    return PreActResNet(
        PreActBlock,
        [1, 1, 1, 1],
        n_cls=n_cls,
        model_width=model_width,
        cuda=cuda,
        half_prec=half_prec,
        activation=activation,
        droprate=droprate,
        bn_flag=bn_flag,
    )


def TinyResNetGroupNorm(
    n_cls, model_width=64, cuda=True, half_prec=False, activation="relu", droprate=0.0
):
    bn_flag = False
    return PreActResNet(
        PreActBlock,
        [1, 1, 1, 1],
        n_cls=n_cls,
        model_width=model_width,
        cuda=cuda,
        half_prec=half_prec,
        activation=activation,
        droprate=droprate,
        bn_flag=bn_flag,
    )


def ResNet18(n_cls, model_width=64):
    return ResNet(
        BasicBlockResNet34, [2, 2, 2, 2], num_classes=n_cls, model_width=model_width
    )


def PreActResNet18(
    n_cls,
    model_width=64,
    cuda=True,
    half_prec=False,
    activation="relu",
    droprate=0.0,
    bn_flag=True,
    learnable_bn=False,
    freeze_last_layer=False,
):
    return PreActResNet(
        PreActBlock,
        [2, 2, 2, 2],
        n_cls=n_cls,
        model_width=model_width,
        cuda=cuda,
        half_prec=half_prec,
        activation=activation,
        droprate=droprate,
        bn_flag=bn_flag,
        learnable_bn=learnable_bn,
        freeze_last_layer=freeze_last_layer,
    )


def PreActResNet34(
    n_cls,
    model_width=64,
    cuda=True,
    half_prec=False,
    activation="relu",
    droprate=0.0,
    bn_flag=True,
    learnable_bn=False,
    freeze_last_layer=False,
):
    return PreActResNet(
        PreActBlock,
        [3, 4, 6, 3],
        n_cls=n_cls,
        model_width=model_width,
        cuda=cuda,
        half_prec=half_prec,
        activation=activation,
        droprate=droprate,
        bn_flag=bn_flag,
        learnable_bn=learnable_bn,
        freeze_last_layer=freeze_last_layer,
    )


def PreActResNet18GroupNorm(
    n_cls, model_width=64, cuda=True, half_prec=False, activation="relu", droprate=0.0
):
    bn_flag = False  # bn_flag==False means that we use GroupNorm with 32 groups
    return PreActResNet(
        PreActBlock,
        [2, 2, 2, 2],
        n_cls=n_cls,
        model_width=model_width,
        cuda=cuda,
        half_prec=half_prec,
        activation=activation,
        droprate=droprate,
        bn_flag=bn_flag,
    )


def PreActResNet34GroupNorm(
    n_cls, model_width=64, cuda=True, half_prec=False, activation="relu", droprate=0.0
):
    bn_flag = False  # bn_flag==False means that we use GroupNorm with 32 groups
    return PreActResNet(
        PreActBlock,
        [3, 4, 6, 3],
        n_cls=n_cls,
        model_width=model_width,
        cuda=cuda,
        half_prec=half_prec,
        activation=activation,
        droprate=droprate,
        bn_flag=bn_flag,
    )


def ResNet34(n_cls, model_width=64):
    return ResNet(
        BasicBlockResNet34, [3, 4, 6, 3], num_classes=n_cls, model_width=model_width
    )


def WideResNet28(n_cls, model_width=10):
    return WideResNet(num_classes=n_cls, widen_factor=model_width)


def get_model(
    model_name,
    n_cls,
    half_prec,
    shapes_dict,
    model_width,
    activation="relu",
    droprate=0.0,
    batch_norm=False,
    freeze_last_layer=False,
    learnable_bn=False,
    bias=True,
):
    from vit_pytorch import ViT, SimpleViT

    if model_name == "vit_basic":
        model = ViT(
            image_size=32,
            patch_size=4,
            num_classes=n_cls,
            dim=model_width,
            depth=6,
            heads=16,
            mlp_dim=model_width * 2,
            dropout=0.1,
            emb_dropout=0.1,
        )
    elif model_name == "vit_exp":
        model = SimpleViT(
            image_size=32,
            patch_size=4,
            num_classes=n_cls,
            dim=model_width,
            depth=6,
            heads=16,
            mlp_dim=model_width * 2,
        )
    elif model_name == "resnet18":
        model = PreActResNet18(
            n_cls,
            model_width=model_width,
            half_prec=half_prec,
            activation=activation,
            droprate=droprate,
            bn_flag=batch_norm,
            freeze_last_layer=freeze_last_layer,
            learnable_bn=learnable_bn,
        )
    elif model_name == "resnet18_plain":
        model = ResNet18(n_cls, model_width)
    elif model_name == "resnet18_gn":
        model = PreActResNet18GroupNorm(
            n_cls,
            model_width=model_width,
            half_prec=half_prec,
            activation=activation,
            droprate=droprate,
        )
    elif model_name == "vgg16":
        assert droprate == 0.0, "dropout is not implemented for vgg16"
        model = VGG16(n_cls, model_width, half_prec, batch_norm, bias)
    elif model_name in ["resnet34", "resnet34_plain"]:
        model = ResNet34(n_cls, model_width)
    elif model_name == "resnet34_gn":
        model = PreActResNet34GroupNorm(
            n_cls,
            model_width=model_width,
            half_prec=half_prec,
            activation=activation,
            droprate=droprate,
        )
    elif model_name == "resnet34preact":
        model = PreActResNet34(
            n_cls,
            model_width=model_width,
            half_prec=half_prec,
            activation=activation,
            droprate=droprate,
            bn_flag=batch_norm,
            freeze_last_layer=freeze_last_layer,
            learnable_bn=learnable_bn,
        )
    elif model_name == "wrn28":
        model = WideResNet28(n_cls, model_width)
    elif model_name == "resnet_tiny":
        model = TinyResNet(
            n_cls,
            model_width=model_width,
            half_prec=half_prec,
            activation=activation,
            droprate=droprate,
        )
    elif model_name == "resnet_tiny_gn":
        model = TinyResNetGroupNorm(
            n_cls,
            model_width=model_width,
            half_prec=half_prec,
            activation=activation,
            droprate=droprate,
        )
    elif model_name == "linear":
        model = Linear(n_cls, shapes_dict)
    else:
        raise ValueError("wrong model")
    model.half_prec = half_prec
    return model


def init_weights(model, scale_init=0.0):
    def init_weights_linear(m):
        if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):
            # m.weight.data.zero_()
            m.weight.data.normal_()
            m.weight.data *= scale_init / (m.weight.data**2).sum() ** 0.5
            if m.bias is not None:
                m.bias.data.zero_()

    def init_weights_he(m):
        # if isinstance(m, nn.Conv2d):
        #     n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
        #     m.weight.data.normal_(0, math.sqrt(2. / n))
        #     if m.bias is not None:
        #         m.bias.data.zero_()
        # elif isinstance(m, nn.Linear):
        #     n = m.in_features
        #     m.weight.data.normal_(0, math.sqrt(2. / n))
        #     if m.bias is not None:
        #         m.bias.data.zero_()

        # From Rice et al.
        if isinstance(m, nn.Conv2d):
            nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
        elif isinstance(m, nn.BatchNorm2d) or isinstance(m, nn.GroupNorm):
            m.weight.data.fill_(1)
            m.bias.data.zero_()
        elif isinstance(m, nn.Linear):
            m.bias.data.zero_()

    if model == "linear":
        return init_weights_linear
    else:
        return init_weights_he


def forward_pass_rlat(model, x, deltas, layers):
    i = 0

    def out_hook(m, inp, out_layer):
        nonlocal i
        if layers[i] == model.normalize:
            new_out = (torch.clamp(inp[0] + deltas[i], 0, 1) - model.mu) / model.std
        else:
            new_out = out_layer + deltas[i]
        i += 1
        return new_out

    handles = [layer.register_forward_hook(out_hook) for layer in layers]
    out = model(x)

    for handle in handles:
        handle.remove()
    return out


def get_rlat_layers(model, layers):
    # import ipdb;ipdb.set_trace()
    if layers == "all":
        return [
            model.normalize,
            model.conv1,
            model.layer1[0].bn1,
            model.layer1[0].conv1,
            model.layer1[0].bn2,
            model.layer1[0].conv2,
            model.layer1[1].bn1,
            model.layer1[1].conv1,
            model.layer1[1].bn2,
            model.layer1[1].conv2,
            model.layer1,
            model.layer2[0].bn1,
            model.layer2[0].conv1,
            model.layer2[0].bn2,
            model.layer2[0].conv2,
            model.layer2[1].bn1,
            model.layer2[1].conv1,
            model.layer2[1].bn2,
            model.layer2[1].conv2,
            model.layer2,
            model.layer3[0].bn1,
            model.layer3[0].conv1,
            model.layer3[0].bn2,
            model.layer3[0].conv2,
            model.layer3[1].bn1,
            model.layer3[1].conv1,
            model.layer3[1].bn2,
            model.layer3[1].conv2,
            model.layer3,
            model.layer4[0].bn1,
            model.layer4[0].conv1,
            model.layer4[0].bn2,
            model.layer4[0].conv2,
            model.layer4[1].bn1,
            model.layer4[1].conv1,
            model.layer4[1].bn2,
            model.layer4[1].conv2,
            model.layer4,
            model.bn,
        ]
    elif layers == "lpips":
        return [model.conv1, model.layer1, model.layer2, model.layer3, model.layer4]
    elif layers == "bnonly":
        return [
            model.normalize,
            model.layer1[0].bn1,
            model.layer1[0].bn2,
            model.layer1[1].bn1,
            model.layer1[1].bn2,
            model.layer2[0].bn1,
            model.layer2[0].bn2,
            model.layer2[1].bn1,
            model.layer2[1].bn2,
            model.layer3[0].bn1,
            model.layer3[0].bn2,
            model.layer3[1].bn1,
            model.layer3[1].bn2,
            model.layer4[0].bn1,
            model.layer4[0].bn2,
            model.layer4[1].bn1,
            model.layer4[1].bn2,
            model.bn,
        ]
    elif layers == "convonly":
        return [
            model.normalize,
            model.conv1,
            model.layer1[0].conv1,
            model.layer1[0].conv2,
            model.layer1[1].conv1,
            model.layer1[1].conv2,
            model.layer2[0].conv1,
            model.layer2[0].conv2,
            model.layer2[1].conv1,
            model.layer2[1].conv2,
            model.layer3[0].conv1,
            model.layer3[0].conv2,
            model.layer3[1].conv1,
            model.layer3[1].conv2,
            model.layer4[0].conv1,
            model.layer4[0].conv2,
            model.layer4[1].conv1,
            model.layer4[1].conv2,
        ]
    elif layers == "block0":
        return [model.conv1]
    elif layers == "block1":
        return [model.layer1]
    elif layers == "block2":
        return [model.layer2]
    elif layers == "block3":
        return [model.layer3]
    elif layers == "block4":
        return [model.layer4]
    else:
        raise ValueError("wrong RLAT layers")


def save_model(epoch, model, optim, log_dir, name):
    torch.save(
        {
            "epoch": epoch,
            "model_state_dict": model.state_dict(),
            "optimizer_state_dict": optim.state_dict(),
        },
        os.path.join(log_dir, "models", f"{name}_{epoch}.pth"),
    )


def _check_bn(module, flag):
    if issubclass(module.__class__, torch.nn.modules.batchnorm._BatchNorm):
        flag[0] = True


def check_bn(model):
    flag = [False]
    model.apply(lambda module: _check_bn(module, flag))
    return flag[0]


def reset_bn(module):
    if issubclass(module.__class__, torch.nn.modules.batchnorm._BatchNorm):
        module.running_mean = torch.zeros_like(module.running_mean)
        module.running_var = torch.ones_like(module.running_var)


def _get_momenta(module, momenta):
    if issubclass(module.__class__, torch.nn.modules.batchnorm._BatchNorm):
        momenta[module] = module.momentum


def _set_momenta(module, momenta):
    if issubclass(module.__class__, torch.nn.modules.batchnorm._BatchNorm):
        module.momentum = momenta[module]


def bn_update(loader, model):
    """
    BatchNorm buffers update (if any).
    Performs 1 epochs to estimate buffers average using train dataset.

    :param loader: train dataset loader for buffers average estimation.
    :param model: model being update
    :return: None
    """
    if not check_bn(model):
        return
    model.train()
    momenta = {}
    model.apply(reset_bn)
    model.apply(lambda module: _get_momenta(module, momenta))
    n = 0
    for input, _ in loader:
        # input = input.cuda(async=True)
        input_var = torch.autograd.Variable(input)
        b = input_var.data.size(0)

        momentum = b / (n + b)
        for module in momenta.keys():
            module.momentum = momentum

        model(input_var)
        n += b

    model.apply(lambda module: _set_momenta(module, momenta))

overparameterized_nets/models/get_models.py
tml-epfl/why-weight-decay
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import math
import os


class Flatten(nn.Module):
    def forward(self, x):
        return x.view(x.size(0), -1)


class Normalize(nn.Module):
    def __init__(self, mu, std):
        super(Normalize, self).__init__()
        self.mu, self.std = mu, std

    def forward(self, x):
        return (x - self.mu) / self.std


class CustomReLU(nn.Module):
    def __init__(self):
        super(CustomReLU, self).__init__()
        self.collect_preact = True
        self.avg_preacts = []

    def forward(self, preact):
        if self.collect_preact:
            self.avg_preacts.append(preact.abs().mean().item())
        act = F.relu(preact)
        return act


class ModuleWithStats(nn.Module):
    def __init__(self):
        super(ModuleWithStats, self).__init__()

    def forward(self, x):
        for layer in self._model:
            if type(layer) == CustomReLU:
                layer.avg_preacts = []

        out = self._model(x)

        avg_preacts_all = [
            layer.avg_preacts for layer in self._model if type(layer) == CustomReLU
        ]
        self.avg_preact = np.mean(avg_preacts_all)
        return out


class Linear(ModuleWithStats):
    def __init__(self, n_cls, shape_in):
        n_cls = 1 if n_cls == 2 else n_cls
        super(Linear, self).__init__()
        d = int(np.prod(shape_in[1:]))
        self._model = nn.Sequential(Flatten(), nn.Linear(d, n_cls, bias=False))

    def forward(self, x):
        logits = self._model(x)
        return torch.cat([torch.zeros(logits.shape).cuda(), logits], dim=1)


class LinearTwoOutputs(ModuleWithStats):
    def __init__(self, n_cls, shape_in):
        super(LinearTwoOutputs, self).__init__()
        d = int(np.prod(shape_in[1:]))
        self._model = nn.Sequential(Flatten(), nn.Linear(d, n_cls, bias=False))


class IdentityLayer(nn.Module):
    def forward(self, inputs):
        return inputs


class PreActBlock(nn.Module):
    """Pre-activation version of the BasicBlock."""

    expansion = 1

    def __init__(
        self,
        in_planes,
        planes,
        bn,
        learnable_bn,
        stride=1,
        activation="relu",
        droprate=0.0,
        gn_groups=32,
    ):
        super(PreActBlock, self).__init__()
        self.collect_preact = True
        self.activation = activation
        self.droprate = droprate
        self.avg_preacts = []
        self.bn1 = (
            nn.BatchNorm2d(in_planes, affine=learnable_bn)
            if bn
            else nn.GroupNorm(gn_groups, in_planes)
        )
        self.conv1 = nn.Conv2d(
            in_planes,
            planes,
            kernel_size=3,
            stride=stride,
            padding=1,
            bias=not learnable_bn,
        )
        self.bn2 = (
            nn.BatchNorm2d(planes, affine=learnable_bn)
            if bn
            else nn.GroupNorm(gn_groups, planes)
        )
        self.conv2 = nn.Conv2d(
            planes, planes, kernel_size=3, stride=1, padding=1, bias=not learnable_bn
        )

        if stride != 1 or in_planes != self.expansion * planes:
            self.shortcut = nn.Sequential(
                # We added this line to ensure scale invariance
                nn.BatchNorm2d(in_planes, affine=learnable_bn),
                nn.Conv2d(
                    in_planes,
                    self.expansion * planes,
                    kernel_size=1,
                    stride=stride,
                    bias=not learnable_bn,
                ),
            )

    def act_function(self, preact):
        if self.activation == "relu":
            act = F.relu(preact)
            # print((act == 0).float().mean().item(), (act.norm() / act.shape[0]).item(), (act.norm() / np.prod(act.shape)).item())
        else:
            assert self.activation[:8] == "softplus"
            beta = int(self.activation.split("softplus")[1])
            act = F.softplus(preact, beta=beta)
        return act

    def forward(self, x):
        out = self.act_function(self.bn1(x))
        shortcut = (
            self.shortcut(out) if hasattr(self, "shortcut") else x
        )  # Important: using out instead of x
        out = self.conv1(out)
        out = self.act_function(self.bn2(out))
        if self.droprate > 0:
            out = F.dropout(out, p=self.droprate, training=self.training)
        out = self.conv2(out)
        out += shortcut
        return out


class BasicBlock(nn.Module):
    def __init__(self, in_planes, out_planes, stride, droprate=0.0):
        super(BasicBlock, self).__init__()
        self.bn1 = nn.BatchNorm2d(in_planes)
        self.relu1 = nn.ReLU(inplace=True)
        self.conv1 = nn.Conv2d(
            in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False
        )
        self.bn2 = nn.BatchNorm2d(out_planes)
        self.relu2 = nn.ReLU(inplace=True)
        self.conv2 = nn.Conv2d(
            out_planes, out_planes, kernel_size=3, stride=1, padding=1, bias=False
        )
        self.droprate = droprate
        self.equalInOut = in_planes == out_planes
        self.convShortcut = (
            (not self.equalInOut)
            and nn.Conv2d(
                in_planes,
                out_planes,
                kernel_size=1,
                stride=stride,
                padding=0,
                bias=False,
            )
            or None
        )

    def forward(self, x):
        if not self.equalInOut:
            x = self.relu1(self.bn1(x))
        else:
            out = self.relu1(self.bn1(x))
        out = self.relu2(self.bn2(self.conv1(out if self.equalInOut else x)))
        if self.droprate > 0:
            out = F.dropout(out, p=self.droprate, training=self.training)
        out = self.conv2(out)
        return torch.add(x if self.equalInOut else self.convShortcut(x), out)


class BasicBlockResNet34(nn.Module):
    expansion = 1

    def __init__(self, in_planes, planes, stride=1):
        super(BasicBlockResNet34, self).__init__()
        self.conv1 = nn.Conv2d(
            in_planes, planes, kernel_size=3, stride=stride, padding=1, bias=False
        )
        self.bn1 = nn.BatchNorm2d(planes)
        self.conv2 = nn.Conv2d(
            planes, planes, kernel_size=3, stride=1, padding=1, bias=False
        )
        self.bn2 = nn.BatchNorm2d(planes)

        self.shortcut = nn.Sequential()
        if stride != 1 or in_planes != self.expansion * planes:
            self.shortcut = nn.Sequential(
                nn.Conv2d(
                    in_planes,
                    self.expansion * planes,
                    kernel_size=1,
                    stride=stride,
                    bias=False,
                ),
                nn.BatchNorm2d(self.expansion * planes),
            )

    def forward(self, x):
        out = F.relu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))
        out += self.shortcut(x)
        out = F.relu(out)
        return out


class NetworkBlock(nn.Module):
    def __init__(self, nb_layers, in_planes, out_planes, block, stride, droprate=0.0):
        super(NetworkBlock, self).__init__()
        self.layer = self._make_layer(
            block, in_planes, out_planes, nb_layers, stride, droprate
        )

    def _make_layer(self, block, in_planes, out_planes, nb_layers, stride, droprate):
        layers = []
        for i in range(int(nb_layers)):
            layers.append(
                block(
                    i == 0 and in_planes or out_planes,
                    out_planes,
                    i == 0 and stride or 1,
                    droprate,
                )
            )
        return nn.Sequential(*layers)

    def forward(self, x):
        return self.layer(x)


class ResNet(nn.Module):
    def __init__(self, block, num_blocks, num_classes=10, model_width=64, droprate=0.0):
        super(ResNet, self).__init__()
        self.in_planes = model_width
        self.half_prec = False
        # self.mu = torch.tensor((0.4914, 0.4822, 0.4465)).view(1, 3, 1, 1).cuda()
        # self.std = torch.tensor((0.2471, 0.2435, 0.2616)).view(1, 3, 1, 1).cuda()
        self.mu = torch.tensor((0.0, 0.0, 0.0)).view(1, 3, 1, 1).cuda()
        self.std = torch.tensor((1.0, 1.0, 1.0)).view(1, 3, 1, 1).cuda()
        # if self.half_prec:
        #     self.mu, self.std = self.mu.half(), self.std.half()

        self.normalize = Normalize(self.mu, self.std)
        self.conv1 = nn.Conv2d(
            3, model_width, kernel_size=3, stride=1, padding=1, bias=False
        )
        self.bn1 = nn.BatchNorm2d(model_width)
        self.layer1 = self._make_layer(block, model_width, num_blocks[0], stride=1)
        self.layer2 = self._make_layer(block, 2 * model_width, num_blocks[1], stride=2)
        self.layer3 = self._make_layer(block, 4 * model_width, num_blocks[2], stride=2)
        self.layer4 = self._make_layer(block, 8 * model_width, num_blocks[3], stride=2)
        self.linear = nn.Linear(8 * model_width * block.expansion, num_classes)

    def _make_layer(self, block, planes, num_blocks, stride):
        strides = [stride] + [1] * (num_blocks - 1)
        layers = []
        for stride in strides:
            layers.append(block(self.in_planes, planes, stride))
            self.in_planes = planes * block.expansion
        return nn.Sequential(*layers)

    def forward(self, x, return_features=False, return_block=5):
        assert return_block in [1, 2, 3, 4, 5], "wrong return_block"
        # out = self.normalize(x)
        out = F.relu(self.bn1(self.conv1(x)))
        out = self.layer1(out)
        if return_features and return_block == 1:
            return out
        out = self.layer2(out)
        if return_features and return_block == 2:
            return out
        out = self.layer3(out)
        if return_features and return_block == 3:
            return out
        out = self.layer4(out)
        if return_features and return_block == 4:
            return out
        out = F.avg_pool2d(out, 4)
        out = out.view(out.size(0), -1)
        if return_features and return_block == 5:
            return out
        out = self.linear(out)
        return out


class PreActResNet(nn.Module):
    def __init__(
        self,
        block,
        num_blocks,
        n_cls,
        model_width=64,
        cuda=True,
        half_prec=False,
        activation="relu",
        droprate=0.0,
        bn_flag=True,
        learnable_bn=True,
        freeze_last_layer=False,
    ):
        super(PreActResNet, self).__init__()
        self.half_prec = half_prec
        self.bn_flag = bn_flag
        # in particular, 32 for model_width=64 as in the original GroupNorm paper
        self.gn_groups = model_width // 2
        self.learnable_bn = learnable_bn  # doesn't matter if self.bn=False
        self.freeze_last_layer = freeze_last_layer
        self.in_planes = model_width
        self.avg_preact = None
        self.activation = activation
        self.n_cls = n_cls
        # self.mu = torch.tensor((0.4914, 0.4822, 0.4465)).view(1, 3, 1, 1)
        # self.std = torch.tensor((0.2471, 0.2435, 0.2616)).view(1, 3, 1, 1)
        self.mu = torch.tensor((0.0, 0.0, 0.0)).view(1, 3, 1, 1)
        self.std = torch.tensor((1.0, 1.0, 1.0)).view(1, 3, 1, 1)

        if cuda:
            self.mu, self.std = self.mu.cuda(), self.std.cuda()
        # if half_prec:
        #     self.mu, self.std = self.mu.half(), self.std.half()

        self.normalize = Normalize(self.mu, self.std)
        self.bn1 = (
            nn.BatchNorm2d(3, affine=self.learnable_bn)
            if self.bn_flag
            else nn.GroupNorm(self.gn_groups, 3)
        )
        self.conv1 = nn.Conv2d(
            3,
            model_width,
            kernel_size=3,
            stride=1,
            padding=1,
            bias=not self.learnable_bn,
        )
        self.layer1 = self._make_layer(block, model_width, num_blocks[0], 1, droprate)
        self.layer2 = self._make_layer(
            block, 2 * model_width, num_blocks[1], 2, droprate
        )
        self.layer3 = self._make_layer(
            block, 4 * model_width, num_blocks[2], 2, droprate
        )
        final_layer_factor = 8
        self.layer4 = self._make_layer(
            block, final_layer_factor * model_width, num_blocks[3], 2, droprate
        )
        self.bn = (
            nn.BatchNorm2d(
                final_layer_factor * model_width * block.expansion,
                affine=self.learnable_bn,
            )
            if self.bn_flag
            else nn.GroupNorm(
                self.gn_groups, final_layer_factor * model_width * block.expansion
            )
        )
        self.linear = nn.Linear(
            final_layer_factor * model_width * block.expansion,
            1 if n_cls == 2 else n_cls,
        )

        if self.freeze_last_layer:
            # We freeze the last linear layer to ensure scale invariance
            self.linear.weight.requires_grad = False
            # We freeze the last linear layer to ensure scale invariance
            self.linear.bias.requires_grad = False

    def _make_layer(self, block, planes, num_blocks, stride, droprate):
        strides = [stride] + [1] * (num_blocks - 1)
        layers = []
        for stride in strides:
            layers.append(
                block(
                    self.in_planes,
                    planes,
                    self.bn_flag,
                    self.learnable_bn,
                    stride,
                    self.activation,
                    droprate,
                    self.gn_groups,
                )
            )
            # layers.append(block(self.in_planes, planes, stride))
            self.in_planes = planes * block.expansion
        return nn.Sequential(*layers)

    def forward(self, x, return_features=False, return_block=5):
        assert return_block in [1, 2, 3, 4, 5], "wrong return_block"
        for layer in [*self.layer1, *self.layer2, *self.layer3, *self.layer4]:
            layer.avg_preacts = []

        # x = x / ((x**2).sum([1, 2, 3], keepdims=True)**0.5 + 1e-6)  # numerical stability is needed for RLAT
        out = self.normalize(x)
        # We added this following Arora's paper about exp learning rate
        out = self.bn1(out)
        out = self.conv1(out)
        out = self.layer1(out)
        out = self.layer2(out)
        out_3 = self.layer3(out)
        out_4 = self.layer4(out_3)
        out_4 = F.relu(self.bn(out_4))
        out_5 = F.avg_pool2d(out_4, 4)
        out_5 = out_5.view(out_5.size(0), -1)

        out = self.linear(out_5)
        if out.shape[1] == 1:
            out = torch.cat([torch.zeros_like(out), out], dim=1)
        if return_features:
            return out, out_3, out_4, out_5
        else:
            return out

    def moving_average(self, model, ema_coef):
        for param1, param2 in zip(self.parameters(), model.parameters()):
            param1.data *= ema_coef
            param1.data += param2.data * (1 - ema_coef)


class WideResNet(nn.Module):
    """Based on code from https://github.com/yaodongyu/TRADES"""

    def __init__(
        self, depth=28, num_classes=10, widen_factor=10, droprate=0.0, bias_last=True
    ):
        super(WideResNet, self).__init__()
        self.half_prec = False
        nChannels = [16, 16 * widen_factor, 32 * widen_factor, 64 * widen_factor]
        assert (depth - 4) % 6 == 0
        n = (depth - 4) / 6
        block = BasicBlock
        # 1st conv before any network block
        self.conv1 = nn.Conv2d(
            3, nChannels[0], kernel_size=3, stride=1, padding=1, bias=False
        )
        # 1st block
        self.block1 = NetworkBlock(n, nChannels[0], nChannels[1], block, 1, droprate)
        # 2nd block
        self.block2 = NetworkBlock(n, nChannels[1], nChannels[2], block, 2, droprate)
        # 3rd block
        self.block3 = NetworkBlock(n, nChannels[2], nChannels[3], block, 2, droprate)
        # global average pooling and classifier
        self.bn1 = nn.BatchNorm2d(nChannels[3])
        self.relu = nn.ReLU(inplace=True)
        self.fc = nn.Linear(nChannels[3], num_classes, bias=bias_last)
        self.nChannels = nChannels[3]

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2.0 / n))
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()
            elif isinstance(m, nn.Linear) and not m.bias is None:
                m.bias.data.zero_()

    def forward(self, x):
        out = self.conv1(x)
        out = self.block1(out)
        out = self.block2(out)
        out = self.block3(out)
        out = self.relu(self.bn1(out))
        out = F.avg_pool2d(out, 8)
        out = out.view(-1, self.nChannels)
        return self.fc(out)


class VGG(nn.Module):
    """
    VGG model. Source: https://github.com/chengyangfu/pytorch-vgg-cifar10/blob/master/vgg.py
    (in turn modified from https://github.com/pytorch/vision.git)
    """

    def __init__(self, n_cls, half_prec, cfg, batch_norm, bias):
        super(VGG, self).__init__()
        self.half_prec = half_prec
        self.mu = torch.tensor((0.485, 0.456, 0.406)).view(1, 3, 1, 1).cuda()
        self.std = torch.tensor((0.229, 0.224, 0.225)).view(1, 3, 1, 1).cuda()
        self.normalize = Normalize(self.mu, self.std)
        self.features = self.make_layers(cfg, batch_norm=batch_norm, bias=bias)
        n_out = cfg[-2]  # equal to 8*model_width
        self.classifier = nn.Sequential(
            # nn.Dropout(),
            nn.Linear(n_out, n_out, bias=bias),
            nn.ReLU(True),
            # nn.Dropout(),
            nn.Linear(n_out, n_out, bias=bias),
            nn.ReLU(True),
            nn.Linear(n_out, n_cls, bias=bias),
        )
        # Initialize weights
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2.0 / n))
                print(bias)
                if bias:
                    m.bias.data.zero_()

    def make_layers(self, cfg, batch_norm=False, bias=False):
        layers = []
        in_channels = 3
        for v in cfg:
            if v == "M":
                layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
            else:
                conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1, bias=bias)
                if batch_norm:
                    layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)]
                else:
                    layers += [conv2d, nn.ReLU(inplace=True)]
                in_channels = v
        return nn.Sequential(*layers)

    def forward(self, x):
        x = self.normalize(x)
        x = self.features(x)
        x = x.view(x.size(0), -1)
        x = self.classifier(x)
        return x

    def moving_average(self, model, ema_coef):
        for param1, param2 in zip(self.parameters(), model.parameters()):
            param1.data *= ema_coef
            param1.data += param2.data * (1 - ema_coef)


def VGG16(n_cls, model_width, half_prec, batch_norm, bias):
    """VGG 16-layer model (configuration "D")"""
    w1, w2, w3, w4, w5 = (
        model_width,
        2 * model_width,
        4 * model_width,
        8 * model_width,
        8 * model_width,
    )
    # cfg = [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M']
    cfg = [w1, w1, "M", w2, w2, "M", w3, w3, w3, "M", w4, w4, w4, "M", w5, w5, w5, "M"]
    return VGG(n_cls, half_prec, cfg, batch_norm, bias)


def TinyResNet(
    n_cls, model_width=64, cuda=True, half_prec=False, activation="relu", droprate=0.0
):
    bn_flag = True
    return PreActResNet(
        PreActBlock,
        [1, 1, 1, 1],
        n_cls=n_cls,
        model_width=model_width,
        cuda=cuda,
        half_prec=half_prec,
        activation=activation,
        droprate=droprate,
        bn_flag=bn_flag,
    )


def TinyResNetGroupNorm(
    n_cls, model_width=64, cuda=True, half_prec=False, activation="relu", droprate=0.0
):
    bn_flag = False
    return PreActResNet(
        PreActBlock,
        [1, 1, 1, 1],
        n_cls=n_cls,
        model_width=model_width,
        cuda=cuda,
        half_prec=half_prec,
        activation=activation,
        droprate=droprate,
        bn_flag=bn_flag,
    )


def ResNet18(n_cls, model_width=64):
    return ResNet(
        BasicBlockResNet34, [2, 2, 2, 2], num_classes=n_cls, model_width=model_width
    )


def PreActResNet18(
    n_cls,
    model_width=64,
    cuda=True,
    half_prec=False,
    activation="relu",
    droprate=0.0,
    bn_flag=True,
    learnable_bn=False,
    freeze_last_layer=False,
):
    return PreActResNet(
        PreActBlock,
        [2, 2, 2, 2],
        n_cls=n_cls,
        model_width=model_width,
        cuda=cuda,
        half_prec=half_prec,
        activation=activation,
        droprate=droprate,
        bn_flag=bn_flag,
        learnable_bn=learnable_bn,
        freeze_last_layer=freeze_last_layer,
    )


def PreActResNet34(
    n_cls,
    model_width=64,
    cuda=True,
    half_prec=False,
    activation="relu",
    droprate=0.0,
    bn_flag=True,
    learnable_bn=False,
    freeze_last_layer=False,
):
    return PreActResNet(
        PreActBlock,
        [3, 4, 6, 3],
        n_cls=n_cls,
        model_width=model_width,
        cuda=cuda,
        half_prec=half_prec,
        activation=activation,
        droprate=droprate,
        bn_flag=bn_flag,
        learnable_bn=learnable_bn,
        freeze_last_layer=freeze_last_layer,
    )


def PreActResNet18GroupNorm(
    n_cls, model_width=64, cuda=True, half_prec=False, activation="relu", droprate=0.0
):
    bn_flag = False  # bn_flag==False means that we use GroupNorm with 32 groups
    return PreActResNet(
        PreActBlock,
        [2, 2, 2, 2],
        n_cls=n_cls,
        model_width=model_width,
        cuda=cuda,
        half_prec=half_prec,
        activation=activation,
        droprate=droprate,
        bn_flag=bn_flag,
    )


def PreActResNet34GroupNorm(
    n_cls, model_width=64, cuda=True, half_prec=False, activation="relu", droprate=0.0
):
    bn_flag = False  # bn_flag==False means that we use GroupNorm with 32 groups
    return PreActResNet(
        PreActBlock,
        [3, 4, 6, 3],
        n_cls=n_cls,
        model_width=model_width,
        cuda=cuda,
        half_prec=half_prec,
        activation=activation,
        droprate=droprate,
        bn_flag=bn_flag,
    )


def ResNet34(n_cls, model_width=64):
    return ResNet(
        BasicBlockResNet34, [3, 4, 6, 3], num_classes=n_cls, model_width=model_width
    )


def WideResNet28(n_cls, model_width=10):
    return WideResNet(num_classes=n_cls, widen_factor=model_width)


def get_model(
    model_name,
    n_cls,
    half_prec,
    shapes_dict,
    model_width,
    activation="relu",
    droprate=0.0,
    batch_norm=False,
    freeze_last_layer=True,
    learnable_bn=False,
    bias=True,
):
    from vit_pytorch import ViT, SimpleViT

    if model_name == "vit_basic":
        model = ViT(
            image_size=32,
            patch_size=4,
            num_classes=n_cls,
            dim=model_width,
            depth=6,
            heads=16,
            mlp_dim=model_width * 2,
            dropout=0.1,
            emb_dropout=0.1,
        )
    elif model_name == "vit_exp":
        model = SimpleViT(
            image_size=32,
            patch_size=4,
            num_classes=n_cls,
            dim=model_width,
            depth=6,
            heads=16,
            mlp_dim=model_width * 2,
        )
    elif model_name == "resnet18":
        model = PreActResNet18(
            n_cls,
            model_width=model_width,
            half_prec=half_prec,
            activation=activation,
            droprate=droprate,
            bn_flag=batch_norm,
            freeze_last_layer=freeze_last_layer,
            learnable_bn=learnable_bn,
        )
    elif model_name == "resnet18_plain":
        model = ResNet18(n_cls, model_width)
    elif model_name == "resnet18_gn":
        model = PreActResNet18GroupNorm(
            n_cls,
            model_width=model_width,
            half_prec=half_prec,
            activation=activation,
            droprate=droprate,
        )
    elif model_name == "vgg16":
        assert droprate == 0.0, "dropout is not implemented for vgg16"
        model = VGG16(n_cls, model_width, half_prec, batch_norm, bias)
    elif model_name in ["resnet34", "resnet34_plain"]:
        model = ResNet34(n_cls, model_width)
    elif model_name == "resnet34_gn":
        model = PreActResNet34GroupNorm(
            n_cls,
            model_width=model_width,
            half_prec=half_prec,
            activation=activation,
            droprate=droprate,
        )
    elif model_name == "resnet34preact":
        model = PreActResNet34(
            n_cls,
            model_width=model_width,
            half_prec=half_prec,
            activation=activation,
            droprate=droprate,
            bn_flag=batch_norm,
            freeze_last_layer=freeze_last_layer,
            learnable_bn=learnable_bn,
        )
    elif model_name == "wrn28":
        model = WideResNet28(n_cls, model_width)
    elif model_name == "resnet_tiny":
        model = TinyResNet(
            n_cls,
            model_width=model_width,
            half_prec=half_prec,
            activation=activation,
            droprate=droprate,
        )
    elif model_name == "resnet_tiny_gn":
        model = TinyResNetGroupNorm(
            n_cls,
            model_width=model_width,
            half_prec=half_prec,
            activation=activation,
            droprate=droprate,
        )
    elif model_name == "linear":
        model = Linear(n_cls, shapes_dict)
    else:
        raise ValueError("wrong model")
    model.half_prec = half_prec
    return model


def init_weights(model, scale_init=0.0):
    def init_weights_linear(m):
        if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):
            # m.weight.data.zero_()
            m.weight.data.normal_()
            m.weight.data *= scale_init / (m.weight.data**2).sum() ** 0.5
            if m.bias is not None:
                m.bias.data.zero_()

    def init_weights_he(m):
        # if isinstance(m, nn.Conv2d):
        #     n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
        #     m.weight.data.normal_(0, math.sqrt(2. / n))
        #     if m.bias is not None:
        #         m.bias.data.zero_()
        # elif isinstance(m, nn.Linear):
        #     n = m.in_features
        #     m.weight.data.normal_(0, math.sqrt(2. / n))
        #     if m.bias is not None:
        #         m.bias.data.zero_()

        # From Rice et al.
        if isinstance(m, nn.Conv2d):
            nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
        elif isinstance(m, nn.BatchNorm2d) or isinstance(m, nn.GroupNorm):
            m.weight.data.fill_(1)
            m.bias.data.zero_()
        elif isinstance(m, nn.Linear):
            m.bias.data.zero_()

    if model == "linear":
        return init_weights_linear
    else:
        return init_weights_he


def forward_pass_rlat(model, x, deltas, layers):
    i = 0

    def out_hook(m, inp, out_layer):
        nonlocal i
        if layers[i] == model.normalize:
            new_out = (torch.clamp(inp[0] + deltas[i], 0, 1) - model.mu) / model.std
        else:
            new_out = out_layer + deltas[i]
        i += 1
        return new_out

    handles = [layer.register_forward_hook(out_hook) for layer in layers]
    out = model(x)

    for handle in handles:
        handle.remove()
    return out


def get_rlat_layers(model, layers):
    # import ipdb;ipdb.set_trace()
    if layers == "all":
        return [
            model.normalize,
            model.conv1,
            model.layer1[0].bn1,
            model.layer1[0].conv1,
            model.layer1[0].bn2,
            model.layer1[0].conv2,
            model.layer1[1].bn1,
            model.layer1[1].conv1,
            model.layer1[1].bn2,
            model.layer1[1].conv2,
            model.layer1,
            model.layer2[0].bn1,
            model.layer2[0].conv1,
            model.layer2[0].bn2,
            model.layer2[0].conv2,
            model.layer2[1].bn1,
            model.layer2[1].conv1,
            model.layer2[1].bn2,
            model.layer2[1].conv2,
            model.layer2,
            model.layer3[0].bn1,
            model.layer3[0].conv1,
            model.layer3[0].bn2,
            model.layer3[0].conv2,
            model.layer3[1].bn1,
            model.layer3[1].conv1,
            model.layer3[1].bn2,
            model.layer3[1].conv2,
            model.layer3,
            model.layer4[0].bn1,
            model.layer4[0].conv1,
            model.layer4[0].bn2,
            model.layer4[0].conv2,
            model.layer4[1].bn1,
            model.layer4[1].conv1,
            model.layer4[1].bn2,
            model.layer4[1].conv2,
            model.layer4,
            model.bn,
        ]
    elif layers == "lpips":
        return [model.conv1, model.layer1, model.layer2, model.layer3, model.layer4]
    elif layers == "bnonly":
        return [
            model.normalize,
            model.layer1[0].bn1,
            model.layer1[0].bn2,
            model.layer1[1].bn1,
            model.layer1[1].bn2,
            model.layer2[0].bn1,
            model.layer2[0].bn2,
            model.layer2[1].bn1,
            model.layer2[1].bn2,
            model.layer3[0].bn1,
            model.layer3[0].bn2,
            model.layer3[1].bn1,
            model.layer3[1].bn2,
            model.layer4[0].bn1,
            model.layer4[0].bn2,
            model.layer4[1].bn1,
            model.layer4[1].bn2,
            model.bn,
        ]
    elif layers == "convonly":
        return [
            model.normalize,
            model.conv1,
            model.layer1[0].conv1,
            model.layer1[0].conv2,
            model.layer1[1].conv1,
            model.layer1[1].conv2,
            model.layer2[0].conv1,
            model.layer2[0].conv2,
            model.layer2[1].conv1,
            model.layer2[1].conv2,
            model.layer3[0].conv1,
            model.layer3[0].conv2,
            model.layer3[1].conv1,
            model.layer3[1].conv2,
            model.layer4[0].conv1,
            model.layer4[0].conv2,
            model.layer4[1].conv1,
            model.layer4[1].conv2,
        ]
    elif layers == "block0":
        return [model.conv1]
    elif layers == "block1":
        return [model.layer1]
    elif layers == "block2":
        return [model.layer2]
    elif layers == "block3":
        return [model.layer3]
    elif layers == "block4":
        return [model.layer4]
    else:
        raise ValueError("wrong RLAT layers")


def save_model(epoch, model, optim, log_dir, name):
    torch.save(
        {
            "epoch": epoch,
            "model_state_dict": model.state_dict(),
            "optimizer_state_dict": optim.state_dict(),
        },
        os.path.join(log_dir, "models", f"{name}_{epoch}.pth"),
    )


def _check_bn(module, flag):
    if issubclass(module.__class__, torch.nn.modules.batchnorm._BatchNorm):
        flag[0] = True


def check_bn(model):
    flag = [False]
    model.apply(lambda module: _check_bn(module, flag))
    return flag[0]


def reset_bn(module):
    if issubclass(module.__class__, torch.nn.modules.batchnorm._BatchNorm):
        module.running_mean = torch.zeros_like(module.running_mean)
        module.running_var = torch.ones_like(module.running_var)


def _get_momenta(module, momenta):
    if issubclass(module.__class__, torch.nn.modules.batchnorm._BatchNorm):
        momenta[module] = module.momentum


def _set_momenta(module, momenta):
    if issubclass(module.__class__, torch.nn.modules.batchnorm._BatchNorm):
        module.momentum = momenta[module]


def bn_update(loader, model):
    """
    BatchNorm buffers update (if any).
    Performs 1 epochs to estimate buffers average using train dataset.

    :param loader: train dataset loader for buffers average estimation.
    :param model: model being update
    :return: None
    """
    if not check_bn(model):
        return
    model.train()
    momenta = {}
    model.apply(reset_bn)
    model.apply(lambda module: _get_momenta(module, momenta))
    n = 0
    for input, _ in loader:
        # input = input.cuda(async=True)
        input_var = torch.autograd.Variable(input)
        b = input_var.data.size(0)

        momentum = b / (n + b)
        for module in momenta.keys():
            module.momentum = momentum

        model(input_var)
        n += b

    model.apply(lambda module: _set_momenta(module, momenta))

overparameterized_nets/models/get_models_scale_inv.py
tml-epfl/why-weight-decay
import pdb
import torch
from torch.cuda.amp import GradScaler, autocast
from configs.config import configuration
from torch.nn import CrossEntropyLoss
from torch.optim import SGD
from tqdm import tqdm
from exp_utils.setup_exp import set_exp
from models import get_models, get_models_scale_inv
from data import get_dataset
from math import sqrt
import copy
from pyhessian import hessian 
from functorch import make_functional, vmap, vjp, jvp, jacrev, make_functional_with_buffers
import numpy as np

SUBSET = 500
JAC = False
HESS = True


def compute_jacobian(loader, model):
    jacs = []
    for i, train_subs_x in enumerate(tqdm(loader)):
        train_subs_x = train_subs_x[0]
        fmodel, params, buffers = make_functional_with_buffers(model)

        def predict(params, buffers, x):
            return fmodel(params, buffers, x)
        (compute_batch_jacobian,) = vmap(jacrev(predict, argnums=(0, )),
                                         in_dims=(None, None, 0))(params, buffers, train_subs_x)
        jac = torch.norm(torch.cat([j.detach().cpu().flatten(
            1) for j in compute_batch_jacobian], 1), dim=1)
        jacs.append(jac.detach().cpu())
        compute_batch_jacobian = 0
        jac = 0
    return torch.cat(jacs).mean().item()


def evaluate(model, loaders, loss_fn, context='test'):
    model.eval()
    with torch.no_grad():
        total_correct, total_num, loss = 0., 0., 0.
        for ims, labs in tqdm(loaders[context], leave=False):
            with autocast():
                # + model(torch.fliplr(ims))) / 2.  Test-time augmentation
                out = (model(ims))
                loss += loss_fn(out, labs).detach().cpu().item()
                total_correct += out.argmax(1).eq(labs).sum().cpu().item()
                total_num += ims.shape[0]
    return 1 - total_correct / total_num, loss / total_num

def eval(model, loaders_eval, loss_fn, ep, aim_writer, context='SDE'): 
    test_err, test_loss = evaluate(
        model, loaders_eval, loss_fn, context='test')
    train_err, train_loss = evaluate(
        model, loaders_eval, loss_fn, context='train')
    l2 = sum([(param**2).sum()
             for param in model.parameters()]).detach().cpu().item()
    aim_writer.track({'Test Error': test_err, 'Test Loss': 0.5*config.wd*l2+test_loss,
                     'Test Cross Entropy': test_loss}, context={'subset': context}, step=ep)
    aim_writer.track({'Train Error': train_err, 'Train Loss': 0.5*config.wd*l2+train_loss,
                     'Train Cross Entropy': train_loss, 'L2': sqrt(l2), 'reg': 0.5*config.wd*l2}, context={'subset': context}, step=ep)

def compute_jac(loaders, model, ep, aim_writer, context='SDE'):
    ############## JACOBIAN COMPUTATION ################
    print("============> COMPUTING JACOBIAN <============")
    jac_norm = compute_jacobian(loaders['train'], model)

    aim_writer.track({'J_norm': jac_norm},  context={'subset': context}, step=ep)


def compute_hess(loaders, model, loss, ep, top_eig=False, aim_writer=None, context='SDE'):
    print("============> COMPUTING HESSIAN <============")
    hessian_comp = hessian(
        model, loss, dataloader=loaders['train'], cuda=True)
    trace = np.mean(hessian_comp.trace())
    if top_eig:
        eigenvalues, _ = hessian_comp.eigenvalues()
        aim_writer.track({'trace_H': trace , 'top_eig':  eigenvalues[-1]},  context={'subset': context}, step=ep)

    else:
        aim_writer.track({'trace_H': trace},  context={'subset': context}, step=ep)


def train(model, loaders, loaders_eval, loss_fn, opt, scaler, ep, aim_writer, ema=None, ma=False, p_radius=None):
    model.train()
    if ma:
        model_ma = copy.deepcopy(model).eval()
    for ims, labs in tqdm(loaders['train'], leave=False):
        opt.zero_grad(set_to_none=True)
        with autocast():
            out = model(ims)
            train_loss = loss_fn(out, labs)

        scaler.scale(train_loss).backward()
        scaler.step(opt)
        scaler.update()
        if ema is not None:
            model_ema.moving_average(model, 0.999)
        if ma:
            with torch.no_grad():
                for param1, param2 in zip(model_ma.parameters(), model.parameters()):
                    param1.data += param2.data
    if p_radius is not None:
        with torch.no_grad():
            norm = torch.norm(
                torch.cat([p.reshape(-1) for p in model.parameters() if p.requires_grad]), p=2)
            for param in model.parameters():
                if param.requires_grad:
                    param.div_(norm/p_radius)
    if ma:
        with torch.no_grad():
            for param in model_ma.parameters():
                param /= len(loaders['train'])
        ma_test_err, ma_test_loss = evaluate(model_ma, loaders_eval, loss_fn)
        aim_writer.track({'Test Error': ma_test_err, 'Test Cross Entropy': ma_test_loss}, context={
                         'subset': 'ma'}, step=ep)


# --------------------------------------------------------------------------------
# SETUP EXPERIMENT
# --------------------------------------------------------------------------------
config = configuration()
device, aim_writer = set_exp(config)
if config.dataset == 'cifar100':    
    n_cls = 100 # TODO: take it from dataset
else: 
    n_cls = 10

if config.scale_inv: 
    model = get_models_scale_inv.get_model(config.model, n_cls, config.half_prec, get_dataset.shapes_dict[config.dataset], config.model_width,
                                batch_norm=config.batch_norm, learnable_bn=False).to(memory_format=torch.channels_last).cuda()
else:
    model = get_models.get_model(config.model, n_cls, config.half_prec, get_dataset.shapes_dict[config.dataset], config.model_width,
                                batch_norm=config.batch_norm, freeze_last_layer=False, learnable_bn=True).to(memory_format=torch.channels_last).cuda()

model_ema = copy.deepcopy(model).eval()
#Init on the sphere
if config.radius is not None: 
    with torch.no_grad():
        norm = torch.norm(
            torch.cat([p.reshape(-1) for p in model.parameters() if p.requires_grad]), p=2)
        for param in model.parameters():
            if param.requires_grad:
                param.div_(norm/config.radius)

start_epoch = 0
opt = SGD(model.parameters(), lr=config.lr,
          momentum=config.momentum, weight_decay=config.wd)
if config.lr_flow is not None:
    gamma_lr = config.lr_flow/config.lr
else:
    gamma_lr = config.lr_gamma_decay

scaler = GradScaler()
loss_fn = CrossEntropyLoss()
loaders = get_dataset.create_dataloaders(
    config.dataset, config.no_data_augm, config.batch_size, device)
loaders_eval = get_dataset.create_dataloaders(
    config.dataset, config.no_data_augm, 1000, device)
h_j_loaders = get_dataset.create_dataloaders(
    'cifar10_5k', config.no_data_augm, 50, device)

# --------------------------------------------------------------------------------
# TRAINING
# --------------------------------------------------------------------------------
with tqdm(range(start_epoch, config.epochs), desc=f'Epochs', unit='epoch') as tepoch:
    for ep in tepoch:
        tepoch.set_description(f"Epoch {ep}")
        train(model, loaders, loaders_eval, loss_fn, opt, scaler, ep, aim_writer, ema=model_ema, ma=False)
        eval(model, loaders_eval, loss_fn, ep, aim_writer, context='SDE')
        aim_writer.track({'learning rate': opt.param_groups[0]['lr'], 'lambda': opt.param_groups[0]['weight_decay']}, context={
                         'subset': 'SDE'}, step=ep)
        get_models.bn_update(loaders['train'], model_ema)
        eval(model_ema, loaders_eval, loss_fn, ep, aim_writer, context='ema')
# --------------------------------------------------------------------------------
# GRADIENT FLOW
# --------------------------------------------------------------------------------
        if ep % config.flow_every == 0:
            model2 = copy.deepcopy(model)
            opt_ft = SGD(model2.parameters(), lr=config.lr_flow,
                         momentum=config.momentum, weight_decay=config.wd_flow)  # Check if you need WD here

            with tqdm(range(ep, ep+config.flow_steps), desc=f'Flow Epochs', unit='epoch', leave=False) as tepoch2:
                for ep2 in tepoch2:
                    tepoch2.set_description(f"Flow {ep}")
                    train(model2, loaders, loaders_eval, loss_fn, opt_ft, scaler, ep2, aim_writer, p_radius=None)
                # I am only evaluating the last point of thefolow to make it faster 
                eval(model2, loaders_eval, loss_fn, ep, aim_writer, context=f'flow_{ep}')
            
            get_models.save_model(
                ep, model, opt, config.out_dir, 'ckpt_SDE')
            get_models.save_model(
                ep, model_ema, opt, config.out_dir, 'ckpt_ema')
            get_models.save_model(
                ep, model2, opt, config.out_dir, 'ckpt_flow')
            
            if JAC:
                compute_jac(h_j_loaders, model2, ep=ep, aim_writer=aim_writer, context=f'flow')
            if HESS:
                compute_hess(h_j_loaders, model2, loss_fn, ep=ep, aim_writer=aim_writer, context=f'flow')
overparameterized_nets/traceh.py
tml-epfl/why-weight-decay
import torch
from torch.cuda.amp import GradScaler, autocast
from configs.config import configuration
from torch.nn import CrossEntropyLoss
from torch.optim import SGD
from tqdm import tqdm
from exp_utils.setup_exp import set_exp
from exp_utils.utils import CustomMultiStepLR
from models import get_models,get_models_scale_inv
from data import get_dataset
from math import sqrt
import copy

def evaluate(model, loaders, loss_fn, context='test'):
    model.eval()
    with torch.no_grad():
        total_correct, total_num, loss = 0., 0., 0.
        for ims, labs in tqdm(loaders[context], leave=False):
            with autocast():
                out = (model(ims))
                loss += loss_fn(out, labs).detach().cpu().item()*ims.shape[0]
                total_correct += out.argmax(1).eq(labs).sum().cpu().item()
                total_num += ims.shape[0]
    return 1 - total_correct / total_num, loss / total_num


def train(model, loaders, loaders_eval, loss_fn, opt, scaler, ep, aim_writer, context, ema=None, ma=False, p_radius=None):
    model.train()
    if ma:
        model_ma = copy.deepcopy(model).eval()
    for ims, labs in tqdm(loaders['train'], leave=False):
        opt.zero_grad(set_to_none=True)
        with autocast():
            out = model(ims)
            train_loss = loss_fn(out, labs)

        scaler.scale(train_loss).backward()
        scaler.step(opt)
        scaler.update()
        if ema is not None:
            model_ema.moving_average(model, 0.999)
        if ma:
            with torch.no_grad():
                for param1, param2 in zip(model_ma.parameters(), model.parameters()):
                    param1.data += param2.data
    if p_radius is not None:
        with torch.no_grad():
            norm = torch.norm(
                torch.cat([p.reshape(-1) for p in model.parameters() if p.requires_grad]), p=2)
            for param in model.parameters():
                if param.requires_grad:
                    param.div_(norm/p_radius)
    if ma:
        with torch.no_grad():
            for param in model_ma.parameters():
                param /= len(loaders['train'])
        ma_test_err, ma_test_loss = evaluate(model_ma, loaders_eval, loss_fn)
        aim_writer.track({'Test Error': ma_test_err, 'Test Cross Entropy': ma_test_loss}, context={
                         'subset': 'ma'}, step=ep)
    test_err, test_loss = evaluate(
        model, loaders_eval, loss_fn, context='test')
    train_err, train_loss = evaluate(
        model, loaders_eval, loss_fn, context='train')
    l2 = sum([(param**2).sum()
             for param in model.parameters()]).detach().cpu().item()
    aim_writer.track({'Test Error': test_err, 'Test Loss': 0.5*config.wd*l2+test_loss,
                     'Test Cross Entropy': test_loss}, context={'subset': context}, step=ep)
    aim_writer.track({'Train Error': train_err, 'Train Loss': 0.5*config.wd*l2+train_loss,
                     'Train Cross Entropy': train_loss, 'L2': sqrt(l2), 'reg': 0.5*config.wd*l2}, context={'subset': context}, step=ep)
    return test_err, train_err, 0.5*config.wd*l2+test_loss, 0.5*config.wd*l2+train_loss, test_loss, train_loss, sqrt(l2)


# --------------------------------------------------------------------------------
# SETUP EXPERIMENT
# --------------------------------------------------------------------------------
config = configuration()
device, aim_writer = set_exp(config)
if config.dataset == 'cifar100':    
    n_cls = 100
else: 
    n_cls = 10

if config.scale_inv: 
    model = get_models_scale_inv.get_model(config.model, n_cls, config.half_prec, get_dataset.shapes_dict[config.dataset], config.model_width,
                                batch_norm=config.batch_norm, learnable_bn=False).to(memory_format=torch.channels_last).cuda()
else:
    model = get_models.get_model(config.model, n_cls, config.half_prec, get_dataset.shapes_dict[config.dataset], config.model_width,
                                batch_norm=config.batch_norm, freeze_last_layer=False, learnable_bn=True).to(memory_format=torch.channels_last).cuda()


model_ema = copy.deepcopy(model).eval()
l2 = sum([(param**2).sum()
         for param in model.parameters()]).detach().cpu().item()
print(l2)
start_epoch = 0
opt = SGD(model.parameters(), lr=config.lr,
          momentum=config.momentum, weight_decay=config.wd)
if config.lr_flow is not None:
    gamma_lr = config.lr_flow/config.lr
else:
    gamma_lr = config.lr_gamma_decay

scheduler = CustomMultiStepLR(opt, [int(
    config.epochs*config.first_decay)], gamma_lr=gamma_lr, gamma_wd=config.wd_gamma_decay)
scaler = GradScaler()
loss_fn = CrossEntropyLoss()
loaders = get_dataset.create_dataloaders(
    config.dataset, config.no_data_augm, config.batch_size, device)
loaders_eval = get_dataset.create_dataloaders(
    config.dataset, config.no_data_augm, 1000, device)

with tqdm(range(start_epoch, config.epochs), desc=f'Epochs', unit='epoch') as tepoch:
    for ep in tepoch:
        tepoch.set_description(f"Epoch {ep}")
        test_err, train_err, test_loss, train_loss, test_CE, train_CE, l2_sde = train(
            model, loaders, loaders_eval, loss_fn, opt, scaler, ep, aim_writer, context='SDE', ema=model_ema, ma=False, p_radius = config.radius)
        scheduler.step()
        aim_writer.track({'learning rate': opt.param_groups[0]['lr'], 'lambda': opt.param_groups[0]['weight_decay']}, context={
                         'subset': 'SDE'}, step=ep)
        get_models.bn_update(loaders['train'], model_ema)
        ema_train_err, ema_train_loss = evaluate(
            model_ema, loaders_eval, loss_fn, context='train')
        ema_test_err, ema_test_loss = evaluate(
            model_ema, loaders_eval, loss_fn, context='test')
        model_ema.eval()
        ema_l2 = sqrt(sum([(param**2).sum()
                      for param in model_ema.parameters()]).detach().cpu().item())

        aim_writer.track({'Test Error': ema_test_err, 'Test Cross Entropy': ema_test_loss, 'Train Error': ema_train_err, 'Train Cross Entropy': ema_train_loss, 'L2': ema_l2}, context={
                         'subset': 'ema'}, step=ep)

        tepoch.set_postfix(Train_CE=train_CE,
                           Train_e=train_err, Test_e=test_err)

overparameterized_nets/train.py
tml-epfl/why-weight-decay