code snippets for [Question] Memory consumption issue. #202

test.py

# NOTE
# 1. 5 dropouts inside the network
# 2. The `raw_data` feeding into the `ax.complete_trial` here are 10 accs + 1 average acc + 1 loss, I actually use 4 accs, 4 F1s, 4 unweithged F1s, 3 averages (of acc, F1 and unweighted F1) and 1 loss instead.

import torch
import numpy as np

from ax.service.ax_client import AxClient
from ax.plot.contour import plot_contour
from ax.plot.trace import optimization_trace_single_method
from ax.service.managed_loop import optimize
from ax.utils.tutorials.cnn_utils import load_mnist
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from sklearn.metrics import accuracy_score, f1_score
from sklearn.utils.sparsefuncs import count_nonzero

torch.manual_seed(12345)
dtype = torch.float
# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
device = torch.device("cpu")

BATCH_SIZE = 512
train_loader, valid_loader, test_loader = load_mnist(batch_size=BATCH_SIZE)


# 5 dropouts inside the network
class CNN(nn.Module):
  """
    Convolutional Neural Network.
    """

  def __init__(self, parameters):
    super().__init__()
    self.conv1 = nn.Conv2d(1, 20, kernel_size=5, stride=1)
    self.fc1 = nn.Linear(8 * 8 * 20, 64)
    self.fc2 = nn.Linear(64, 10)

  def forward(self, x):
    x = F.relu(self.conv1(x))
    x = F.max_pool2d(x, 3, 3)
    # x = F.dropout(
    #     x, p=parameters.get("d1", 0.001), training=self.training, inplace=False)
    # x = F.dropout(
    #     x, p=parameters.get("d2", 0.001), training=self.training, inplace=False)
    x = x.view(-1, 8 * 8 * 20)
    # x = F.dropout(
    #     x, p=parameters.get("d3", 0.001), training=self.training, inplace=False)
    # x = F.dropout(
    #     x, p=parameters.get("d4", 0.001), training=self.training, inplace=False)
    x = F.relu(self.fc1(x))
    # x = F.dropout(
    #     x, p=parameters.get("d5", 0.001), training=self.training, inplace=False)
    x = self.fc2(x)
    return F.log_softmax(x, dim=-1)


def train(net, train_loader, dtype, device):
  """
    Train CNN on provided data set.

    Args:
        net: initialized neural network
        train_loader: DataLoader containing training set
        parameters: dictionary containing parameters to be passed to the optimizer.
            - lr: default (0.001)
            - momentum: default (0.0)
            - weight_decay: default (0.0)
            - num_epochs: default (1)
        dtype: torch dtype
        device: torch device
    Returns:
        nn.Module: trained CNN.
    """
  # Initialize network
  net.to(dtype=dtype, device=device)  # pyre-ignore [28]
  net.train()
  # Define loss and optimizer
  criterion = nn.NLLLoss(reduction="sum")
  optimizer = optim.SGD(
      net.parameters(),
      lr=0.001,
      momentum=0.0,
      weight_decay=0.0,
  )
  scheduler = optim.lr_scheduler.StepLR(
      optimizer,
      step_size=int(parameters.get("step_size", 30)),
      gamma=parameters.get("gamma", 1.0),  # default is no learning rate decay
  )
  num_epochs = parameters.get("num_epochs", 1)

  # Train Network
  for _ in range(num_epochs):
    for inputs, labels in train_loader:
      # move data to proper dtype and device
      inputs = inputs.to(dtype=dtype, device=device)
      labels = labels.to(device=device)

      # zero the parameter gradients
      optimizer.zero_grad()

      # forward + backward + optimize
      outputs = net(inputs)
      loss = criterion(outputs, labels)
      loss.backward()
      optimizer.step()
      scheduler.step()

  return net


def evaluate(net, data_loader, dtype, device):
  """
    Compute classification accuracy on provided dataset.

    Args:
        net: trained model
        data_loader: DataLoader containing the evaluation set
        dtype: torch dtype
        device: torch device
    Returns:
        float: classification accuracy
    """
  net.eval()
  correct = 0
  total = 0
  cost = 0
  criterion = nn.NLLLoss(reduction="sum")
  with torch.no_grad():
    for inputs, labels in data_loader:
      # move data to proper dtype and device
      inputs = inputs.to(dtype=dtype, device=device)
      labels = labels.to(device=device)
      outputs = net(inputs)
      loss = criterion(outputs, labels)
    _, predicted = torch.max(outputs.data, 1)
    predicted = predicted.numpy()
    labels = labels.numpy()
    score = labels == predicted
    correct += np.array([score[labels == i].sum() for i in range(10)])
    total += np.array([labels[labels == i].size for i in range(10)])
    cost += loss.item()
  return correct / total, correct.sum() / total.sum(), cost / total.sum()


def train_evaluate(parameterization):
  net = CNN(parameterization)
  net = train(net=net, train_loader=train_loader, dtype=dtype, device=device)
  result = evaluate(
      net=net,
      data_loader=valid_loader,
      dtype=dtype,
      device=device,
  )

  # here are 10 accs + 1 average acc + 1 loss, I actually use 4 accs, 4 F1s, 4 unweithged F1s, 3 averages (of acc, F1 and unweighted F1) and 1 loss instead.
  raw_data = {f"acc{idx}": value for idx, value in enumerate(result[0])}
  raw_data.update({"accAvg": result[1]})
  raw_data.update({"loss": result[2]})
  return raw_data


parameters = [{
    "name": "d1",
    "type": "range",
    "bounds": [0, 1.],
    "value_type": "float",
    "log_scale": False
}, {
    "name": "d2",
    "type": "range",
    "bounds": [0, 1.],
    "value_type": "float",
    "log_scale": False
}, {
    "name": "d3",
    "type": "range",
    "bounds": [0, 1.],
    "value_type": "float",
    "log_scale": False
}, {
    "name": "d4",
    "type": "range",
    "bounds": [0, 1.],
    "value_type": "float",
    "log_scale": False
}, {
    "name": "d5",
    "type": "range",
    "bounds": [0, 1.],
    "value_type": "float",
    "log_scale": False
}]

ax = AxClient()
ax.create_experiment(
    name='tmp',
    parameters=parameters,
    objective_name="accAvg",
    minimize=False,
    parameter_constraints=None,
    outcome_constraints=None)
while (True):
  parameters, trial_index = ax.get_next_trial()
  raw_data = train_evaluate(parameters)
  raw_data = {key: (value, 0.) for key, value in raw_data.items()}
  ax.complete_trial(trial_index=trial_index, raw_data=raw_data)