Weights & Biasesを使用したデータサイエンス実験管理

#🐝　wandbのインストール
!pip install wandb -Uq
...
# 必要なライブラリ・モジュールのインストール・インポート
...
#🐝　wandのインポートとログイン
import wandb
wandb.login()
# project nameの指定
project_name = "mnist-pytorch-demo"

def load_and_log():
  #🐝　wandbのrunを作成
  with wandb.init(project=project_name, job_type="load-data") as run:
      datasets = load()  # separate code for loading the datasets
      names = ["training", "validation", "test"]

      #🐝 Artifactを定義
      raw_data = wandb.Artifact(
          "mnist-raw", type="dataset",
          description="Raw MNIST dataset, split into train/val/test",
          metadata={"source": "torchvision.datasets.MNIST","sizes": [len(dataset) for dataset in datasets]}
          )

      for name, subset in zip(names, datasets):
          # Artifactへの保存のためにsubdatasetをTensorDatasetに変換
          data = transform_subset2tensordataset(subset)
          with raw_data.new_file(name + ".pt", mode="wb") as file:
              x, y = data.tensors
              torch.save((x, y), file)

      #🐝 ArtifactをW&Bに保存
      run.log_artifact(raw_data)
...
load_and_log()

def preprocess_and_log(steps):
  #🐝　wandbのinitialization
  with wandb.init(project=project_name, job_type="preprocess-data") as run:
      processed_data = wandb.Artifact(
          "mnist-preprocess", type="dataset",
          description="Preprocessed MNIST dataset",
          metadata=steps)

      #🐝 どのArtifactを使用するかを定義し、必要に応じてダウンロード
      raw_data_artifact = run.use_artifact('mnist-raw:v0')
      raw_dataset = raw_data_artifact.download()

      for split in ["training", "validation", "test"]:
          raw_split = read(raw_dataset, split)
          processed_dataset = preprocess(raw_split, **steps)
          with processed_data.new_file(split + ".pt", mode="wb") as file:
              x, y = processed_dataset.tensors
              torch.save((x, y), file)

      run.log_artifact(processed_data)

table = wandb.Table(columns=["image", "label","split"])

config = dict(
    epochs=10,
    kernels=[16, 32],
    batch_size=128,
    dropout =0.1,
    learning_rate=0.005,
    optimizer="adam",
    log_prediction_iamge=True,
    classes=10,
    dataset="MNIST",
    architecture="CNN"
    )
	
def model_pipeline(config):
    #🐝　wandbのrunを作成
    with wandb.init(project=project_name, config=config) as run:
      config = wandb.config
      # モデルアーキテクチャの定義・データローダー・損失関数・最適化手法の設定
      model, train_loader, validation_loader, criterion, optimizer = make(config, run)
      # モデルをトレーニングし、検証データに対して検証
      train(model, train_loader,validation_loader, criterion, optimizer, config, run)
　　　#🐝　Artifactにモデルを保存
      artifact = wandb.Artifact('model', type='model')
      with  artifact.new_file('model.pt',mode='wb') as file:
          torch.save(model.state_dict(), file)
      run.log_artifact(artifact)
    return model
...
def train(model, training_loader, validation_loader, criterion, optimizer, config, run):
    #🐝　wandb.watchを用いてモデルのgradientsやweightsなどをトラック
    wandb.watch(model, criterion, log="all", log_freq=10)
    # 学習の開始
    example_ct = 0
    for epoch in tqdm(range(config.epochs)):
        for step, (images, labels) in enumerate(training_loader):
            loss = train_batch(images, labels, model, optimizer, criterion)
            example_ct +=  len(images)
            #🐝　run.logを用いてモデルのlossをトラック
            run.log({"epoch": epoch, "train_loss": loss}, step=example_ct, commit=False)
        validation(model, validation_loader, criterion,  run, log_images=config.log_prediction_iamge)
...
model_pipeline(config)


sweep_config = {
    'method': 'random',
    'metric': {
        'name': 'validation_loss',
        'goal': 'minimize'
        },
    }
parameters_dict = {
    'optimizer': {
        'values': ['adam', 'sgd']
        },
    'dropout': {
          'values': [0.1, 0.2, 0.3]
        },
    'learning_rate': {
        # 0から0.1の間を一様分布で乱数を発生
        'distribution': 'uniform',
        'min': 0,
        'max': 0.1
      },
    'batch_size': {
        # 32から256の間で対数スケールにおいて等間隔の分布から乱数を発生
        'distribution': 'q_log_uniform_values',
        'q': 8,
        'min': 32,
        'max': 256,
      },
    'kernels':{
        'values':[[16,32],[32,64]]
      },
    'epochs':{
        'value':5
    },
    'classes':{
        'value':10,
    },
    'dataset':{
        'value':'MNIST',
    },
    'architecture':{
        'value':'CNN',
    },
    'log_prediction_iamge':{
        'value':False
    },
    }

sweep_config['parameters'] = parameters_dict

#🐝 Sweep用に関数を作成
def model_pipeline_sweep(config=None):
    #🐝　wandbのrunを作成
    with wandb.init(config=config) as run:
      config = wandb.config
      # モデルアーキテクチャの定義・データローダー・損失関数・最適化手法の設定
      model, train_loader, validation_loader, criterion, optimizer = make(config, run)
      # モデルをトレーニングし、検証データに対して検証
      train(model, train_loader, validation_loader, criterion, optimizer, config, run)
　　　#🐝　Artifactにモデルを保存
      artifact = wandb.Artifact('model', type='model')
      with  artifact.new_file('model.pt',mode='wb') as file:
          torch.save(model.state_dict(), file)
      run.log_artifact(artifact)

sweep_id = wandb.sweep(sweep_config, project=project_name)
wandb.agent(sweep_id=sweep_id, function=model_pipeline_sweep, count=5)

# sweepの後であれば、sweepのagentとの連携をまず切ります
wandb.sdk.wandb_setup._setup(_reset=True)
# もっとも精度の良いモデルのconfigを取得します
import pandas as pd
import numpy as np
api = wandb.Api()
# Replace 'username' and 'project' with your username and project name
runs = api.runs(project_name)
name_list, path_list, sweepName_list,validation_loss_list = [], [], [], []
for run in runs:
    name_list.append(run.name)
    path_list.append(run.path)
    sweepName_list.append(run.sweepName)
    validation_loss_list.append(run.summary.get("validation_loss", np.nan))

runs_df = pd.DataFrame({
    "name": name_list,
    "path": path_list,
    "sweepName": sweepName_list,
    "validation_loss": validation_loss_list
    })

best_sweep_path = runs_df.query("sweepName.notnull() and validation_loss == validation_loss.min()")["path"].values[0]
best_sweep_path = '/'.join(best_sweep_path)
best_sweep_path

def test_pipeline(config):
    #🐝　wandbのrunを作成
    with wandb.init(project=project_name, config=config) as run:
      config = wandb.config
      # モデルアーキテクチャの定義・データローダー・損失関数・最適化手法の設定
      model, train_loader, validation_loader,test_loader, criterion, optimizer = make(config, run)
      # モデルをトレーニングし、検証データに対して検証
      train(model, train_loader,validation_loader, criterion, optimizer, config, run)
      # テストデータで検証
      test(model, test_loader, criterion, run, config)

      #🐝　Artifactにモデルを保存
      artifact = wandb.Artifact('model', type='model')
      with  artifact.new_file('model.pt',mode='wb') as file:
          torch.save(model.state_dict(), file)
      run.log_artifact(artifact)

def test(model, test_loader, criterion, run, config):
    model.eval()
    table = wandb.Table(columns=["image", "pred", "target"]+[f"score_{i}" for i in range(10)]) # wandbのTableを作成
    with torch.no_grad():
        correct, total, test_loss = 0, 0, 0
        for images, labels in test_loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
            test_loss += criterion(outputs, labels)*labels.size(0)
            if config.log_prediction_iamge:
                for img, pred, targ, prob in zip(images.to("cpu"), predicted.to("cpu"), labels.to("cpu"), outputs.softmax(dim=1).to("cpu")):
                  table.add_data(wandb.Image(img[0].numpy()*255), pred, targ, *prob.numpy())
                run.log({"test(holdout)_prediction_table":table}, commit=False)
        run.log({"test_loss(holdout)":test_loss / total ,"test_accuracy(holdout)": correct / total})

test_pipeline(final_config)