### Start Training (SGD)

Source: https://github.com/microsoft/mup/blob/main/examples/Transformer/README.md

Scale up the model and start training using the SGD optimizer with μP. Ensure base shapes are loaded for proper scaling.

```bash
python main.py --d_model 4096 --load_base_shapes width256.bsh --optimizer musgd --lr 0.5 --cuda
```

--------------------------------

### Install mup from Source

Source: https://github.com/microsoft/mup/blob/main/README.md

Install the mup library from its source code. This involves cloning the repository, navigating to the directory, and installing dependencies and the package itself.

```bash
pip install -r requirements.txt
pip install -e .
```

--------------------------------

### Install mup Package

Source: https://github.com/microsoft/mup/blob/main/README.md

Install the mup library using pip. This is the standard way to get the library for use in your projects.

```bash
pip install mup
```

--------------------------------

### Start Training (Adam)

Source: https://github.com/microsoft/mup/blob/main/examples/Transformer/README.md

Scale up the model and start training using the Adam optimizer with μP. Ensure base shapes are loaded for proper scaling.

```bash
python main.py --d_model 4096 --load_base_shapes width256.bsh --optimizer muadam --lr 0.01 --cuda
```

--------------------------------

### Start Training μP ResNet Model

Source: https://github.com/microsoft/mup/blob/main/examples/ResNet/README.md

Scale up the model and train using the same hyperparameters as the small model. This command initiates training for wider models with μP optimization.

```bash
# for SGD
python main.py --width_mult 2 --optimizer musgd
```

```bash
# for Adam
python main.py --width_mult 2 --optimizer muadam
```

--------------------------------

### Instantiate and Set Base Shapes for Model Scaling

Source: https://github.com/microsoft/mup/blob/main/README.md

Demonstrates how to instantiate base, delta, and target models, then use `set_base_shapes` to configure mup for parameter scaling. This should be called early in the model setup process.

```python
### Instantiate a base model
base_model = MyModel(width=1)
### Optionally, use `torchdistx.deferred_init.deferred_init` to avoid instantiating the parameters
### Simply install `torchdistx` and use
# base_model = torchdistx.deferred_init.deferred_init(MyModel, width=1)
### Instantiate a "delta" model that differs from the base model
###   in all dimensions ("widths") that one wishes to scale.
### Here it's simple, but e.g., in a Transformer, you may want to scale
###   both nhead and dhead, so the delta model should differ in both.
delta_model = MyModel(width=2) # Optionally use `torchdistx` to avoid instantiating

### Instantiate the target model (the model you actually want to train).
### This should be the same as the base model except 
###   the widths could be potentially different.
### In particular, base_model and model should have the same depth.
model = MyModel(width=100)

### Set base shapes
### When `model` has same parameter shapes as `base_model`,
###   `model` behaves exactly the same as `base_model`
###   (which is in PyTorch's default parametrization).
###   This provides backward compatibility at this particular model size.
###   Otherwise, `model`'s init and LR are scaled by μP.
### IMPORTANT: this should be called as soon as possible,
###   before re-initialization and optimizer definition.
set_base_shapes(model, base_model, delta=delta_model)

### Alternatively, one can save the base model shapes in a file
# make_base_shapes(base_model, delta_model, filename)
### and later set base shapes directly from the filename
# set_base_shapes(model, filename)
### This is useful when one cannot fit both 
###   base_model and model in memory at the same time
```

--------------------------------

### Start MLP Training

Source: https://github.com/microsoft/mup/blob/main/examples/MLP/README.md

Begin training a μP MLP model using previously saved base shapes. If --load_base_shapes is omitted, a SP model will be trained by default.

```python
python main.py --load_base_shapes width64.bsh
```

--------------------------------

### Run Mup Tests

Source: https://github.com/microsoft/mup/blob/main/README.md

Execute the test suite for the Mup project using Python's module execution. Ensure Python is installed and accessible in your environment.

```bash
python -m mup.test
```

--------------------------------

### MuP Training Output Log

Source: https://github.com/microsoft/mup/blob/main/examples/MLP/demo.ipynb

Example output log from the muP MLP training process, showing training progress per epoch, loss values, and processing time per batch.

```text
Output:
loading base shapes from ./demo_width256.bsh
done
Train Epoch: 1 [0/50000 (0%)]	Loss: 2.302585 | ms/batch  0.20
Train Epoch: 1 [19200/50000 (38%)]	Loss: 1.897947 | ms/batch  4.87
Train Epoch: 1 [38400/50000 (77%)]	Loss: 1.781119 | ms/batch  4.79

Train set: Average loss: 1.9597

Train Epoch: 2 [0/50000 (0%)]	Loss: 2.022595 | ms/batch  0.19
Train Epoch: 2 [19200/50000 (38%)]	Loss: 1.872523 | ms/batch  5.09
Train Epoch: 2 [38400/50000 (77%)]	Loss: 1.772421 | ms/batch  4.74

Train set: Average loss: 1.7464

Train Epoch: 3 [0/50000 (0%)]	Loss: 1.573457 | ms/batch  0.18
Train Epoch: 3 [19200/50000 (38%)]	Loss: 1.533341 | ms/batch  4.82
Train Epoch: 3 [38400/50000 (77%)]	Loss: 1.798549 | ms/batch  4.93

Train set: Average loss: 1.6913

Train Epoch: 4 [0/50000 (0%)]	Loss: 1.718129 | ms/batch  0.20
Train Epoch: 4 [19200/50000 (38%)]	Loss: 1.719680 | ms/batch  4.75
Train Epoch: 4 [38400/50000 (77%)]	Loss: 1.562585 | ms/batch  4.67

Train set: Average loss: 1.6538

Train Epoch: 5 [0/50000 (0%)]	Loss: 1.684843 | ms/batch  0.18
Train Epoch: 5 [19200/50000 (38%)]	Loss: 1.695738 | ms/batch  4.58
Train Epoch: 5 [38400/50000 (77%)]	Loss: 1.630646 | ms/batch  4.62

Train set: Average loss: 1.6212

Train Epoch: 6 [0/50000 (0%)]	Loss: 1.618959 | ms/batch  0.18
Train Epoch: 6 [19200/50000 (38%)]	Loss: 1.663781 | ms/batch  4.60
Train Epoch: 6 [38400/50000 (77%)]	Loss: 1.839550 | ms/batch  4.76

Train set: Average loss: 1.5935

Train Epoch: 7 [0/50000 (0%)]	Loss: 1.425978 | ms/batch  0.18
Train Epoch: 7 [19200/50000 (38%)]	Loss: 1.581953 | ms/batch  4.83
Train Epoch: 7 [38400/50000 (77%)]	Loss: 1.490716 | ms/batch  4.60

Train set: Average loss: 1.5682

Train Epoch: 8 [0/50000 (0%)]	Loss: 1.578557 | ms/batch  0.17
Train Epoch: 8 [19200/50000 (38%)]	Loss: 1.448391 | ms/batch  4.74
Train Epoch: 8 [38400/50000 (77%)]	Loss: 1.436007 | ms/batch  4.76

Train set: Average loss: 1.5439

Train Epoch: 9 [0/50000 (0%)]	Loss: 1.262302 | ms/batch  0.21
Train Epoch: 9 [19200/50000 (38%)]	Loss: 1.450943 | ms/batch  4.65
Train Epoch: 9 [38400/50000 (77%)]	Loss: 1.502101 | ms/batch  4.69

Train set: Average loss: 1.5228

Train Epoch: 10 [0/50000 (0%)]	Loss: 1.409971 | ms/batch  0.16
Train Epoch: 10 [19200/50000 (38%)]	Loss: 1.564316 | ms/batch  4.66
Train Epoch: 10 [38400/50000 (77%)]	Loss: 1.847680 | ms/batch  4.66

Train set: Average loss: 1.5008

Train Epoch: 11 [0/50000 (0%)]	Loss: 1.416700 | ms/batch  0.18
Train Epoch: 11 [19200/50000 (38%)]	Loss: 1.551162 | ms/batch  4.75
Train Epoch: 11 [38400/50000 (77%)]	Loss: 1.344760 | ms/batch  4.74

Train set: Average loss: 1.4807

Train Epoch: 12 [0/50000 (0%)]	Loss: 1.540586 | ms/batch  0.18
Train Epoch: 12 [19200/50000 (38%)]	Loss: 1.584867 | ms/batch  4.66
Train Epoch: 12 [38400/50000 (77%)]	Loss: 1.520993 | ms/batch  4.61

Train set: Average loss: 1.4620

Train Epoch: 13 [0/50000 (0%)]	Loss: 1.356099 | ms/batch  0.18
Train Epoch: 13 [19200/50000 (38%)]	Loss: 1.287571 | ms/batch  4.73
Train Epoch: 13 [38400/50000 (77%)]	Loss: 1.303671 | ms/batch  4.69

Train set: Average loss: 1.4434

Train Epoch: 14 [0/50000 (0%)]	Loss: 1.549423 | ms/batch  0.18
Train Epoch: 14 [19200/50000 (38%)]	Loss: 1.338795 | ms/batch  4.78
Train Epoch: 14 [38400/50000 (77%)]	Loss: 1.189449 | ms/batch  4.79
```

--------------------------------

### Define Model with MuReadout and Handle Attention Scaling

Source: https://github.com/microsoft/mup/blob/main/README.md

Example of defining a PyTorch model that uses MuReadout for its output layer and adjusts attention scaling for transformer models. This ensures compatibility with mup's scaling mechanisms.

```python
from mup import MuReadout, make_base_shapes, set_base_shapes, MuSGD, MuAdam

class MyModel(nn.Module):
    def __init__(self, width, ...):
        ...
        ### In model definition, replace output layer with MuReadout
        # readout = nn.Linear(width, d_out)
        readout = MuReadout(width, d_out)
        ### If tying weights with an input nn.Embedding layer, do
        # readout = MuSharedReadout(input_layer.weight)
        ...
    def forward(self, ...):
        ...
        ### If using a transformer, make sure to use
        ###   1/d instead of 1/sqrt(d) attention scaling
        # attention_scores = query @ key.T / d**0.5
        attention_scores = query @ key.T * 8 / d
        ### We use 8/d instead of 1/d here to be backward compatible
        ###   with 1/d**0.5 when d=64, a common head dimension.
        ...
```

--------------------------------

### Import necessary libraries for coordinate check

Source: https://github.com/microsoft/mup/blob/main/examples/ResNet/CoordCheck.ipynb

Imports the `coord_check` function and plotting libraries. Ensure these libraries are installed.

```python
from main import coord_check
from itertools import product
import seaborn as sns
import matplotlib.pyplot as plt
sns.set()
```

--------------------------------

### Example Coord Check Plot Generation

Source: https://github.com/microsoft/mup/blob/main/README.md

This snippet demonstrates the typical workflow for generating coordinate check plots using `mup.coord_check`. It involves constructing a dictionary of models with varying widths, preparing a dataloader, recording activation data, and plotting the results. Ensure custom initializations are handled within the lambda function.

```Python
from mup.coord_check import get_coord_data, plot_coord_data
# construct a dictionary of lazy μP models with differing widths
def lazy_model(width):
    # `set_base_shapes` returns the model
    return lambda: set_base_shapes(MyMuModel(width), 'my/base/shape/path.bsh')
    # Note: any custom initialization with `mup.init` would need to
    # be done inside the lambda as well
models = {64: lazy_model(64), ..., 1024: lazy_model(1024)}
# make a dataloader with small batch size/seq len
#   just for testing
dataloader = ...
# record data from the model activations over a few steps of training
# this returns a pandas dataframe
df = get_coord_data(models, dataloader)
# This saves the coord check plots to filename.
plot_coord_data(df, save_to=filename)
# If you are in jupyter notebook, you can also do
#   `plt.show()`
# to show the plot
```

--------------------------------

### Verify Implementation with Coordinate Check (SGD)

Source: https://github.com/microsoft/mup/blob/main/examples/Transformer/README.md

Check the size of activation coordinates as model width increases using SGD optimizer. Load base shapes and enable coordinate check.

```bash
python main.py --load_base_shapes width256.bsh --optimizer sgd --lr 0.5 --cuda --coord_check
```

--------------------------------

### Verify Implementation with Coordinate Check (Adam)

Source: https://github.com/microsoft/mup/blob/main/examples/Transformer/README.md

Check the size of activation coordinates as model width increases using Adam optimizer. Load base shapes and enable coordinate check.

```bash
python main.py --load_base_shapes width256.bsh --optimizer adam --lr 0.01 --cuda --coord_check
```

--------------------------------

### Initialize Device and Random Seed

Source: https://github.com/microsoft/mup/blob/main/examples/MLP/demo.ipynb

Sets the manual seed for PyTorch for reproducibility and moves the computation to the CUDA device if available.

```python
torch.manual_seed(1)
device = torch.device("cuda")
kwargs = {'num_workers': 1, 'pin_memory': True}
```

--------------------------------

### Replace Manual Initialization with mup.init

Source: https://github.com/microsoft/mup/blob/main/README.md

When manually initializing model parameters, replace PyTorch's default initialization functions with their equivalents from `mup.init` to ensure compatibility with mup's scaling.

```python
### Replace your custom init, if any
for param in model.parameters():
    ### If initializing manually with fixed std or bounds, 
    ### then replace with same function from mup.init
    # torch.nn.init.uniform_(param, -0.1, 0.1)
    mup.init.uniform_(param, -0.1, 0.1)
    ### Likewise, if using
    ###   `xavier_uniform_, xavier_normal_, kaiming_uniform_, kaiming_normal_`
    ### from `torch.nn.init`, replace with the same functions from `mup.init`
```

--------------------------------

### Import Libraries and Define Constants

Source: https://github.com/microsoft/mup/blob/main/examples/MLP/demo.ipynb

Imports necessary libraries like PyTorch, NumPy, and torchvision. Defines constants for batch size, epochs, and data directory.

```python
import time
import math

import numpy as np
import torch.nn.functional as F
from torchvision import datasets, transforms
import torch
from torch import nn
from torch.optim import SGD
import seaborn as sns
import pandas as pd

from mup import MuSGD, get_shapes, set_base_shapes, MuReadout


'''
    PyTorch MLP on CIFAR-10, with μP.

    This is minimal notebook that demonstrates the hyperparameter stability of muP.
'''

batch_size = 64
epochs = 20
log_interval = 300
nonlin = torch.relu
criterion = F.cross_entropy
data_dir = '/tmp'
base_shapes_path = './demo_width256.bsh'
```

--------------------------------

### Verify MLP Implementation with Coordinate Check

Source: https://github.com/microsoft/mup/blob/main/examples/MLP/README.md

Check the size of activation coordinates as model width increases to verify the μP implementation. This uses helper functions from mup and generates plots for comparison.

```bash
python main.py --load_base_shapes width64.bsh --coord_check
```

--------------------------------

### Verify ResNet Implementation with Coordinate Check

Source: https://github.com/microsoft/mup/blob/main/examples/ResNet/README.md

Check the size of activation coordinates as model width increases. This is recommended before scaling up and training. Plots are generated to visualize coordinate sizes under μP and SP.

```bash
# for SGD
python main.py --load_base_shapes resnet18.bsh --optimizer sgd --lr 0.1 --coord_check
```

```bash
# for Adam
python main.py --load_base_shapes resnet18.bsh --optimizer adam --lr 0.001 --coord_check
```

--------------------------------

### Save Model Base Shapes

Source: https://github.com/microsoft/mup/blob/main/examples/Transformer/README.md

Specify and save the base shapes for a μP model. This is the first step before training.

```python
python main.py --d_model 256 --save_base_shapes width256.bsh
```

--------------------------------

### Use MuSGD Optimizer

Source: https://github.com/microsoft/mup/blob/main/README.md

Replace standard PyTorch optimizers with their mup counterparts, such as MuSGD, to leverage mup's learning rate scaling capabilities. This should be done after setting base shapes and potentially re-initializing parameters.

```python
### Use the optimizers from `mup.optim` instead of `torch.optim`
# optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
optimizer = MuSGD(model.parameters(), lr=0.1)

### Then just train normally
```

--------------------------------

### Run coordinate check with different configurations

Source: https://github.com/microsoft/mup/blob/main/examples/Transformer/CoordCheck.ipynb

Iterates through different configurations of 'mup' (model-parallel), learning rate, and optimizer to perform coordinate checks. Each check is followed by displaying the plot.

```python
for mup, (lr, optimizer) in product([False, True], [(0.01, 'adam'), (0.3, 'sgd')]):
  args = Namespace(nhead=2, ffn_ratio=1, nlayers=2, dropout=0,
                   tied=False, bias=False, init_var=1, bptt=35,
                  load_base_shapes='width256.bsh', device='cuda',
                  attn_mult=1, output_mult=1, precision='half')
  coord_check(mup=mup,
              lr=lr, optimizer=optimizer, batch_size=32, nsteps=3,
              nseeds=1, data_dir='./data/wikitext-2', args=args,
              plotdir='coord_checks', legend=False)
  plt.show()
```

--------------------------------

### Standard Parametrization (SP) MLP Training Loop

Source: https://github.com/microsoft/mup/blob/main/examples/MLP/demo.ipynb

This snippet shows the training loop for an SP MLP model. It iterates through different model widths and learning rates, configuring the optimizer and running the training for a specified number of epochs. Use this for training MLPs with varying widths and learning rates.

```python
for width in [256, 512, 1024, 2048, 4096, 8192]:
    for log2lr in np.linspace(-8, 0, 20):
        torch.manual_seed(1)
        mynet = MLP(width=width, nonlin=nonlin, output_mult=output_mult, input_mult=input_mult).to(device)
        optimizer = SGD(mynet.parameters(), lr=2**log2lr)
        for epoch in range(1, epochs+1):
            train_loss = train(mynet, device, train_loader, optimizer, epoch, criterion=criterion)
            logs.append(dict(
                epoch=epoch,
                model_type='SP MLP',
                log2lr=log2lr,
                train_loss=train_loss,
                width=width,
            ))
            if math.isnan(train_loss):
                break
```

--------------------------------

### Load CIFAR-10 Dataset

Source: https://github.com/microsoft/mup/blob/main/examples/MLP/demo.ipynb

Loads the CIFAR-10 training and testing datasets using torchvision. Includes transformations for converting images to tensors and normalizing them.

```python
transform = transforms.Compose(
        [transforms.ToTensor(),
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
trainset = datasets.CIFAR10(root=data_dir, train=True,
                                        download=True, transform=transform)
train_loader = torch.utils.data.DataLoader(trainset, batch_size=batch_size,
                                        shuffle=True, num_workers=2)
testset = datasets.CIFAR10(root=data_dir, train=False,
                                    download=True, transform=transform)
test_loader = torch.utils.data.DataLoader(testset, batch_size=batch_size,
                                        shuffle=False, num_workers=2)
```

--------------------------------

### Loading Model Checkpoint with MuMu Base Shapes

Source: https://github.com/microsoft/mup/blob/main/README.md

After loading a PyTorch model checkpoint, manually re-apply base shapes using `set_base_shapes` with `rescale_params=False`. This is necessary because `torch.save` does not preserve mup's `infshape` objects.

```python
model = torch.load('my/model/path.pt')
# Important: note the flag `rescale_params=False`!
set_base_shapes(model, 'my/base/shape/path.bsh', rescale_params=False)
```

--------------------------------

### Import Libraries and Define Plotting Function

Source: https://github.com/microsoft/mup/blob/main/CoordCheck.ipynb

Imports necessary libraries for plotting and defines the `example_plot_coord_check` function. Ensure these imports are present before running the plotting code.

```python
from mup.coord_check import example_plot_coord_check
from itertools import product
import seaborn as sns
import matplotlib.pyplot as plt
sns.set()
```

--------------------------------

### Run coordinate check with different configurations

Source: https://github.com/microsoft/mup/blob/main/examples/ResNet/CoordCheck.ipynb

Iterates through different combinations of MUP (True/False), learning rates, and optimizers to perform coordinate checks. Each check saves a plot to the 'coord_checks' directory and displays it.

```python
for mup, (lr, optimizer) in product([False, True], [(0.01, 'adam'), (0.01, 'sgd')]):
  coord_check(mup=mup,
              lr=lr, optimizer=optimizer, nsteps=3, arch='resnet18',
              base_shapes='resnet18.bsh', nseeds=1, plotdir='coord_checks', legend=False)
  plt.show()
```

--------------------------------

### Define Training Hyperparameters

Source: https://github.com/microsoft/mup/blob/main/examples/MLP/demo.ipynb

Sets optimal hyperparameters for training, specifically the output and input multipliers used in the muP training process.

```python
# optimal HPs
output_mult = 32
input_mult = 0.00390625
```

--------------------------------

### Import necessary libraries

Source: https://github.com/microsoft/mup/blob/main/examples/Transformer/CoordCheck.ipynb

Imports the required libraries for coordinate checking, including the main function, argument parsing, itertools, and plotting libraries.

```python
from main import coord_check
from argparse import Namespace
from itertools import product
import seaborn as sns
import matplotlib.pyplot as plt
sns.set()
```

--------------------------------

### MuP MLP Training Loop

Source: https://github.com/microsoft/mup/blob/main/examples/MLP/demo.ipynb

Trains a muP MLP model using MuAdamW optimizer across different widths and learning rates. Logs training loss and epoch information. Ensure `base_shapes_path` is correctly set.

```python
logs = []

for width in [256, 512, 1024, 2048, 4096, 8192]:
    for log2lr in np.linspace(-8, 0, 20):
        torch.manual_seed(1)
        mynet = muMLP(width=width, nonlin=nonlin, output_mult=output_mult, input_mult=input_mult).to(device)
        print(f'loading base shapes from {base_shapes_path}')
        set_base_shapes(mynet, base_shapes_path)
        print('done')
        optimizer = MuSGD(mynet.parameters(), lr=2**log2lr)
        for epoch in range(1, epochs+1):
            train_loss = train(mynet, device, train_loader, optimizer, epoch, criterion=criterion)
            logs.append(dict(
                epoch=epoch,
                model_type='muP MLP',
                log2lr=log2lr,
                train_loss=train_loss,
                width=width,
            ))
            if math.isnan(train_loss):
                break
```

--------------------------------

### Prepare Base Shapes for MLP

Source: https://github.com/microsoft/mup/blob/main/examples/MLP/demo.ipynb

Defines and saves base shapes for an MLP, specifying which dimensions are widths for scaling during training. Requires the MLP model and its non-linearity function.

```python
base_shapes = get_shapes(MLP(width=256, nonlin=nonlin))
delta_shapes = get_shapes(
    # just need to change whatever dimension(s) we are scaling
    MLP(width=256+1, nonlin=nonlin)
)
make_base_shapes(base_shapes, delta_shapes, savefile=base_shapes_path)
```

--------------------------------

### Relative Learning Rate Adjustment in MuAdam Optimizer

Source: https://github.com/microsoft/mup/blob/main/README.md

When using MuAdam, adjust learning rates relatively within parameter groups. Avoid setting learning rates absolutely, as this can interfere with mup's scaling mechanisms.

```python
optimizer = mup.MuAdam(model.parameters(), lr=1e-3)
for pg in optimizer.param_groups:
  # what NOT to do: setting learning rate absolutely
  # pg['lr'] = 1e-3 * 2
  # what is an OK alternative: setting it relatively
  pg['lr'] *= 2
```

--------------------------------

### Standard MLP Model Definition

Source: https://github.com/microsoft/mup/blob/main/examples/MLP/demo.ipynb

Defines a standard Multi-Layer Perceptron (MLP) model in PyTorch without biases. It includes configurable width, nonlinearity, and multipliers for input and output layers. Parameters are initialized using Kaiming normal distribution.

```python
class MLP(nn.Module):
    def __init__(self, width=128, num_classes=10, nonlin=F.relu, output_mult=1.0, input_mult=1.0):
        super(MLP, self).__init__()
        self.nonlin = nonlin
        self.input_mult = input_mult
        self.output_mult = output_mult
        self.fc_1 = nn.Linear(3072, width, bias=False)
        self.fc_2 = nn.Linear(width, width, bias=False)
        self.fc_3 = nn.Linear(width, num_classes, bias=False)
        self.reset_parameters()
    
    def reset_parameters(self):
        nn.init.kaiming_normal_(self.fc_1.weight, a=1, mode='fan_in')
        self.fc_1.weight.data /= self.input_mult**0.5
        nn.init.kaiming_normal_(self.fc_2.weight, a=1, mode='fan_in')
        nn.init.zeros_(self.fc_3.weight)

    def forward(self, x):
        out = self.nonlin(self.fc_1(x) * self.input_mult**0.5)
        out = self.nonlin(self.fc_2(out))
        return self.fc_3(out) * self.output_mult
```

--------------------------------

### Save Model Base Shapes with mup

Source: https://github.com/microsoft/mup/blob/main/examples/ResNet/README.md

Specify the base shapes for training a μP model. This command saves the base shapes information to a file.

```python
python main.py --save_base_shapes resnet18.bsh --width_mult 1
```

--------------------------------

### Visualize muP MLP Training Loss

Source: https://github.com/microsoft/mup/blob/main/examples/MLP/demo.ipynb

Generates a line plot showing the training loss versus log2 learning rate for the muP MLP model. Filters are applied for model type, training loss less than 5, and epoch 20.

```python
logs_df = pd.DataFrame(logs)
sns.lineplot(x='log2lr', y='train_loss', hue='width', data=logs_df[(logs_df['model_type']=='muP MLP')&
                                                                   (logs_df['train_loss']<5)&
                                                                   (logs_df['epoch']==20)])
```

--------------------------------

### Generate Coordinate Check Plots

Source: https://github.com/microsoft/mup/blob/main/CoordCheck.ipynb

Iterates through various combinations of architecture, optimizer, batch normalization, and mup settings to generate and display coordinate check plots. Plots are saved to the 'coord_checks/' directory.

```python
for arch, opt, bn, mup in product(['mlp', 'cnn'], ['sgd', 'adam'], [False, True], [False, True]):
  example_plot_coord_check(arch, opt, batchnorm=bn, mup=mup, nseeds=5, download_cifar=True, legend=None,
                  plotdir='coord_checks/')
  plt.show()
```

--------------------------------

### PyTorch MLP Training Loop

Source: https://github.com/microsoft/mup/blob/main/examples/MLP/demo.ipynb

Defines a function to train a PyTorch model. It iterates through the training data, computes loss, performs backpropagation, and updates model weights. Includes logging of loss and time per batch.

```python
def train(model, device, train_loader, optimizer, epoch,
            scheduler=None, criterion=F.cross_entropy):
    model.train()
    train_loss = 0
    correct = 0
    start_time = time.time()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data.view(data.size(0), -1))
        
        loss = criterion(output, target)
        loss.backward()
        train_loss += loss.item() * data.shape[0]  # sum up batch loss
        optimizer.step()
        if batch_idx % log_interval == 0:
            elapsed = time.time() - start_time
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f} | ms/batch {:5.2f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                100. * batch_idx / len(train_loader), loss.item(),
                elapsed * 1000 / log_interval))
            start_time = time.time()
        if scheduler is not None:
            scheduler.step()
    train_loss /= len(train_loader.dataset)
    print('\nTrain set: Average loss: {:.4f}\n'.format(
        train_loss, correct, len(train_loader.dataset)))
    return train_loss
```

--------------------------------

### Visualize SP MLP Training Loss

Source: https://github.com/microsoft/mup/blob/main/examples/MLP/demo.ipynb

Generates a line plot illustrating the training loss against log2 learning rate for the SP MLP model. Data is filtered for model type, training loss below 5, and epoch 20.

```python
sns.lineplot(x='log2lr', y='train_loss', hue='width', data=logs_df[(logs_df['model_type']=='SP MLP')&
                                                                   (logs_df['train_loss']<5)&
                                                                   (logs_df['epoch']==20)])
```

--------------------------------

### Create directory for coordinate check plots

Source: https://github.com/microsoft/mup/blob/main/examples/ResNet/CoordCheck.ipynb

Creates a directory named 'coord_checks' if it does not already exist. This is used to store generated plot images.

```python
import os
os.makedirs('coord_checks', exist_ok=True)
```

--------------------------------

### Save MLP Base Shapes

Source: https://github.com/microsoft/mup/blob/main/examples/MLP/README.md

Specify and save the base shapes for the narrowest μP model. This is a prerequisite for training wider models.

```python
python main.py --save_base_shapes width64.bsh
```

--------------------------------

### MLP Shape Representation

Source: https://github.com/microsoft/mup/blob/main/examples/MLP/demo.ipynb

Illustrates the shape representation of an MLP's weights after processing. This is a result of the `get_shapes` function.

```text
Result:
{'fc_1.weight': InfShape([InfDim(256, None), InfDim(None, None)]),
 'fc_2.weight': InfShape([InfDim(256, None), InfDim(256, None)]),
 'fc_3.weight': InfShape([InfDim(None, None), InfDim(256, None)])}
```

--------------------------------

### MuP-Enhanced MLP Model Definition

Source: https://github.com/microsoft/mup/blob/main/examples/MLP/demo.ipynb

Defines an MLP model enhanced with MuP (μP) using the MuReadout layer for the output. It shares a similar structure to the standard MLP but incorporates MuP-specific components for potential hyperparameter stability.

```python
class muMLP(nn.Module):
    def __init__(self, width=128, num_classes=10, nonlin=F.relu, output_mult=1.0, input_mult=1.0):
        super(muMLP, self).__init__()
        self.nonlin = nonlin
        self.input_mult = input_mult
        self.output_mult = output_mult
        self.fc_1 = nn.Linear(3072, width, bias=False)
        self.fc_2 = nn.Linear(width, width, bias=False)
        self.fc_3 = MuReadout(width, num_classes, bias=False, output_mult=self.output_mult)
        self.reset_parameters()
    
    def reset_parameters(self):
        nn.init.kaiming_normal_(self.fc_1.weight, a=1, mode='fan_in')
        self.fc_1.weight.data /= self.input_mult**0.5
        nn.init.kaiming_normal_(self.fc_2.weight, a=1, mode='fan_in')
        nn.init.zeros_(self.fc_3.weight)

    def forward(self, x):
        out = self.nonlin(self.fc_1(x) * self.input_mult**0.5)
        out = self.nonlin(self.fc_2(out))
        return self.fc_3(out)
```

=== COMPLETE CONTENT === This response contains all available snippets from this library. No additional content exists. Do not make further requests.