W&B Training and POC Template

wandb_logger = WandbLogger(project='MNIST', log_model='all', save_code=True, ) # log all new checkpoints during training

training_loader = DataLoader(training_set, batch_size=64, shuffle=True, pin_memory=True)
validation_loader = DataLoader(validation_set, batch_size=64, pin_memory=True)
## Using a raw DataLoader, rather than LightningDataModule, for greater transparency

# Set up model
model = MNIST_LitModule(n_layer_1=128, n_layer_2=128)

trainer = Trainer(gpus=None, max_epochs=5, profiler="pytorch",logger=wandb_logger,
                          callbacks=[
                                     log_predictions_callback, 
                                     checkpoint_callback
                                     ], 
                      precision=32)
trainer.profiler.dirpath="/content/wandb/latest-run/tbprofile"
trainer.fit(model, training_loader, validation_loader)
# trace_files = glob.glob("/content/lightning_logs/*.pt.trace.json")
trace_files = glob.glob("/content/wandb/latest-run/tbprofile/*.pt.trace.json")
for i, trace_file in enumerate(trace_files):
    if "training_step" in trace_file:
      profile_art = wandb.Artifact(f"train-trace{i}-{wandb.run.id}", type="profile")
      profile_art.add_file(trace_file, "train_trace.pt.trace.json")
    else:
      profile_art = wandb.Artifact(f"validation-trace{i}-{wandb.run.id}", type="profile")
      profile_art.add_file(trace_file, "validation_trace.pt.trace.json")
    wandb.log_artifact(profile_art)
wandb.finish()

import wandb

with wandb.init(project="my_project", job_type="model_training") as run:
  # Create Artifact
  training_images = wandb.Artifact(name='training_images', type="training_data")
  # Add serialized data
  training_images.add_dir('sample_images')
  # Log to W&B, automatic versioning
  wandb.log_artifact(training_images)

# Log by reference if data sits in cloud object stores or remote files systems 
with wandb.init(project="my_project", job_type="model_training") as run:
  # Create Artifact
  training_images = wandb.Artifact(name='training_images_reference', type="training_data")
  # Add reference to data, store only the metadata associated with artifact
  training_images.add_reference(uri='file:///content/sample_images')
  # Log to W&B, automatic versioning
  wandb.log_artifact(training_images)

import wandb
run = wandb.init()
artifact = run.use_artifact('tim-w/credit_scorecard/decent-puddle-21_model.json:v0', type='training_data')
artifact_dir = artifact.download()

def train(config=None):
    # Initialize a new wandb run
    with wandb.init(config=config):
        # If called by wandb.agent, as below,
        # this config will be set by Sweep Controller
        config = wandb.config

        loader = build_dataset(config.batch_size)
        network = build_network(config.fc_layer_size, config.dropout)
        optimizer = build_optimizer(network, config.optimizer, config.learning_rate)

        for epoch in range(config.epochs):
            avg_loss = train_epoch(network, loader, optimizer)
            wandb.log({"loss": avg_loss, "epoch": epoch})

sweep_config = {
    'method': 'random',
    # 'method': 'grid',
    # 'method': 'bayes',
    }

parameters_dict = {
    'optimizer': {
        'values': ['adam', 'sgd']
    },
    'fc_layer_size': {
        'values': [128, 256, 512, 1024]
    },
    'dropout': {
        'values': [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7]
    },
    # Static hyperparameter, notice singular key
    'epochs': {
        'value': 10
    },
    'learning_rate': {
        # Flat distribution between 0 and 0.1
        'distribution': 'uniform',
        'min': 0,
        'max': 0.25
    },
    'batch_size': {
        'distribution': 'q_log_uniform',
        'q': 1,
        'min': math.log(32),
        'max': math.log(256),
    }
}

### Initializes the central sweep server
sweep_config['parameters'] = parameters_dict
sweep_id = wandb.sweep(sweep_config, project="sweeps-demo-pytorch")

### Run this in multiple machines/cores to distribute the hyperparameter search
wandb.agent(sweep_id, train, count=10)

import torch
import wandb
from torch.utils.tensorboard import SummaryWriter

sweep_config = {
  'method': 'grid', 
  'metric': {
      'name': 'Loss/train',  ## matches what i write via SummaryWriter
      'goal': 'minimize'
  },
  'early_terminate':{
      'type': 'hyperband',
      'min_iter': 5
  },
  'parameters': {
      'learning_rate':{
          'values': [0.01, 0.005, 0.001, 0.0005, 0.0001]
      }
  }
}

x = torch.arange(-5, 5, 0.1).view(-1, 1)
y = -5 * x + 0.1 * torch.randn(x.size())

wandb.login()
wandb.tensorboard.patch(root_logdir = "/content/runs"). ## since i don't have tensorboard running, i set the root logging directory manually

def sweep_train(config_defaults = None): 
    config_defaults = { "learning_rate": 0.01 }  ## will be overwritten during sweeps
    run = wandb.init(config=config_defaults, sync_tensorboard=True)  # the config gets over-written in the Sweep
    model = torch.nn.Linear(1, 1)
    writer = SummaryWriter(log_dir = "/content/runs")
    criterion = torch.nn.MSELoss()
    optimizer = torch.optim.SGD(model.parameters(), lr = wandb.config.learning_rate)
    for epoch in range(5):
        y1 = model(x)
        loss = criterion(y1, y)
        writer.add_scalar("Loss/train", loss, epoch)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    writer.flush()
    writer.close()
    run.finish()

sweep_id = wandb.sweep(sweep_config, project="sweeps-pytorch-tensorboard-v4")
wandb.agent(sweep_id, function=sweep_train)