<jupyter_start><jupyter_text>Ray Train - A Library for Distributed Deep Learning© 2019-2022, Anyscale. All Rights Reserved![Anyscale Academy](../images/AnyscaleAcademyLogo.png)[Ray Train](https://docs.ray.io/en/latest/train/train.html) is a lightweight library for distributed deep learning. It provides thin wrappers around [PyTorch](https://pytorch.org), [TensorFlow](https://tensorflow.org), and [Horvod](https://horovod.ai/) native modules for data parallel training.> **NOTE**: Ray SGD is renamed to Ray Train  PyTorch Fashion MNIST for Distributed Training We will use Ray Train to distribute our training using couple of models and evaluating which of the two provides usthe best accuracy and a minimal loss. As excercise, you can try to further investigate how you improve the model—via regularization techniques, using CNN layers, trying different loss functions.The steps we will follow are no different (may be slight variation but the essence is the same) from the previous notenbooks.So let's go! First, do the necessary imports, as before.<jupyter_code>import os
from typing import Dict

import torch
import torch.nn.functional as F

import ray
import ray.train as train
from ray.train.trainer import Trainer
from ray.train.callbacks import JsonLoggerCallback, TBXLoggerCallback
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor<jupyter_output><empty_output><jupyter_text>Step 1: Download Train and test datasets<jupyter_code>training_data = datasets.FashionMNIST(
    root="~/data",
    train=True,
    download=True,
    transform=ToTensor(),
)
# Download test data from open datasets.
test_data = datasets.FashionMNIST(
    root="~/data",
    train=False,
    download=True,
    transform=ToTensor(),
)<jupyter_output><empty_output><jupyter_text>Step 2: Define a Neural Network Models. This is a quite simple NN model<jupyter_code># Define model-1
class NeuralNetwork(nn.Module):
    def __init__(self):
        super(NeuralNetwork, self).__init__()
        self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(28 * 28, 512), nn.ReLU(), nn.Linear(512, 512), nn.ReLU(),
            nn.Linear(512, 10), nn.ReLU())

    def forward(self, x):
        x = self.flatten(x)
        logits = self.linear_relu_stack(x)
        return logits<jupyter_output><empty_output><jupyter_text>Define a deeper NN model archiecture with dropouts[source](https://medium.com/@aaysbt/fashion-mnist-data-training-using-pytorch-7f6ad71e96f4)<jupyter_code># Define model-2
class Classifier(nn.Module):
  def __init__(self):
    super().__init__()
    self.fc1 = nn.Linear(784, 120)
    self.fc2 = nn.Linear(120, 120)
    self.fc3 = nn.Linear(120,10)
    self.dropout = nn.Dropout(0.2)

  def forward(self,x):
    x = x.view(x.shape[0],-1)
    x = self.dropout(F.relu(self.fc1(x)))
    x = self.dropout(F.relu(self.fc2(x)))
    x = F.log_softmax(self.fc3(x), dim=1)
    return x
# Define accuracy function
def accuracy_fn(y_pred, y_true):
    n_correct = torch.eq(y_pred, y_true).sum().item()
    acc = (n_correct / len(y_pred)) * 100
    return acc<jupyter_output><empty_output><jupyter_text>Step 3: Define per epoch training and validation functinos<jupyter_code>def train_epoch(dataloader, model, loss_fn, optimizer, epoch):
    size = len(dataloader.dataset)
    for batch, (X, y) in enumerate(dataloader):
        # Compute prediction error
        pred = model(X)
        loss = loss_fn(pred, y)
        
        # Backpropagation
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
def validate_epoch(dataloader, model, loss_fn, epoch):
    size = len(dataloader.dataset)
    num_batches = len(dataloader)
    model.eval()
    test_loss, correct, acc =  0, 0, 0.0
    with torch.no_grad():
        for X, y in dataloader:
            pred = model(X)
            test_loss += loss_fn(pred, y).item()
            correct += (pred.argmax(1) == y).type(torch.float).sum().item()
            predictions = pred.max(dim=1)[1]
            acc += accuracy_fn(predictions, y)
    test_loss /= num_batches
    acc /= num_batches
    correct /= size
    if epoch > 0 and epoch % 50 == 0:
        print(f"Epoc: {epoch}, Avg validation loss: {test_loss:.2f}, Avg validation accuracy: {acc:.2f}%") 
        print("--" * 40)
    return test_loss<jupyter_output><empty_output><jupyter_text>Step 4: Define Ray Train Training functionThis function will be passed to `train.run(...)`<jupyter_code>def train_func(config: Dict):
    batch_size = config.get("batch_size", 64) 
    lr = config.get('lr', 1e-3)
    epochs = config.get("epochs", 20)
    momentum = config.get("momentum", 0.9)
    model_type = config.get('model_type', None)
    loss_fn = config.get("loss_fn", nn.NLLLoss())

    # Create data loaders.
    train_dataloader = DataLoader(training_data, batch_size=batch_size)
    test_dataloader = DataLoader(test_data, batch_size=batch_size)

    # Prepare to use Ray integrated wrappers around PyTorch's Dataloaders
    train_dataloader = train.torch.prepare_data_loader(train_dataloader)
    test_dataloader = train.torch.prepare_data_loader(test_dataloader)

    # Create model.

    model = Classifier() if model_type else NeuralNetwork()
    # Prepare to use Ray integrated wrappers around PyTorch's model
    model = train.torch.prepare_model(model)
    
    # Get or objective loss function
    loss_fn = config.get("loss_fn", nn.NLLLoss())

    optimizer = torch.optim.SGD(model.parameters(), lr=lr, momentum=momentum)

    loss_results = []

    for e in range(epochs):
        train_epoch(train_dataloader, model, loss_fn, optimizer, e)
        loss = validate_epoch(test_dataloader, model, loss_fn, e)
        train.report(loss=loss)
        loss_results.append(loss)

    return loss_results<jupyter_output><empty_output><jupyter_text>Step 5: Wrap our Trainer around a main driver function<jupyter_code>def train_fashion_mnist(num_workers=12, use_gpu=False):
    trainer = Trainer(
        backend="torch", num_workers=num_workers, use_gpu=use_gpu)
    trainer.start()
    result = trainer.run(
        train_func=train_func,
        config={
            "lr": 1e-3,
            "batch_size": 128,
            "epochs": 150,
            "momentum": 0.9,
            "model_type": 0,                     # change to 1 for second NN model
            "loss_fn": nn.CrossEntropyLoss()     # change to nn.nn.NLLLoss() 
        },
        callbacks=[JsonLoggerCallback(), TBXLoggerCallback()])
    trainer.shutdown() 
    return result<jupyter_output><empty_output><jupyter_text>Step 6: Define some parallelism parameters And a URL to connect to a Ray Cluster if running on Anysacle<jupyter_code>number_of_workers = 8
use_gpu = False                              # change to True if using a Ray cluster with GPUs
address = "anyscale://ray_train_ddp_cluster" # use your anyscale cluster here<jupyter_output><empty_output><jupyter_text>Step 6: Connect to Ray cluster<jupyter_code>CONNECT_TO_ANYSCALE=False
if ray.is_initialized:
    ray.shutdown()
    if CONNECT_TO_ANYSCALE:
        ray.init(address)
    else:
        ray.init(ignore_reinit_error=True)<jupyter_output>2022-03-16 16:30:27,154	INFO services.py:1412 -- View the Ray dashboard at [1m[32mhttp://127.0.0.1:8268[39m[22m<jupyter_text>Step 7: Run the main Trainer driver<jupyter_code>%%time
results = train_fashion_mnist(num_workers=number_of_workers, use_gpu=use_gpu)<jupyter_output>2022-03-16 16:19:11,394	INFO trainer.py:199 -- Trainer logs will be logged in: /Users/jules/ray_results/train_2022-03-16_16-19-11
[2m[36m(BaseWorkerMixin pid=61696)[0m 2022-03-16 16:19:13,471	INFO torch.py:66 -- Setting up process group for: env:// [rank=0, world_size=8]
[2m[36m(BaseWorkerMixin pid=61697)[0m 2022-03-16 16:19:13,458	INFO torch.py:66 -- Setting up process group for: env:// [rank=1, world_size=8]
[2m[36m(BaseWorkerMixin pid=61699)[0m 2022-03-16 16:19:13,466	INFO torch.py:66 -- Setting up process group for: env:// [rank=2, world_size=8]
[2m[36m(BaseWorkerMixin pid=61690)[0m 2022-03-16 16:19:13,473	INFO torch.py:66 -- Setting up process group for: env:// [rank=5, world_size=8]
[2m[36m(BaseWorkerMixin pid=61690)[0m [W ProcessGroupGloo.cpp:707] Warning: Unable to resolve hostname to a (local) address. Using the loopback address as fallback. Manually set the network interface to bind to with GLOO_SOCKET_IFNAME. (function operator())
[2m[36m(BaseWorkerMixin pid[...]<jupyter_text>Step 8: Observe metrics in Tensorboard Subsitute your path `train_path` printed in the cell above<jupyter_code>!tensorboard --logdir ~/ray_results/<train_path>
ray.shutdown()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.11451169188445667, 'fraction_numerical': 0.036680421824850984, 'mean_word_length': 3.960204661739625, 'pattern_counts': {'":': 6, '<': 44, '<?xml version=': 0, '>': 46, 'https://': 5, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 1, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '11', 'dir': 'jupyter-structured-clean-dedup', 'id': '166452', 'license': 'Apache-2.0', 'n_tokens_mistral': 3158, 'n_tokens_neox': 2909, 'n_words': 726, 'path': 'ray-train/04-Ray-Train-with-PyTorch-FashionMNIST.ipynb', 'repo_name': 'RomaKoks/academy', 'size': '8724'}