PyTorch + DASK

#Torch packages
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torchvision import datasets
import torch.nn.functional as F
from torchvision import models

#Packages to handle the distribution
from dask.distributed import Client
import torch.distributed as dist
from dask_jobqueue import PBSCluster
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler

#System modules
import os
import time
import socket

PROCESSES = 60
EPOCHS_COUNT = 5

BASE_PATH = "/sciclone/home/dsmillerrunfol/AML/TORCH/"
IMAGE_PATH = "/sciclone/home/dsmillerrunfol/AML/TORCH/UCMerced_LandUse/Images"

def logger(string, path=str(BASE_PATH+"distributedlog.log")):
    string = str(string)
    with open(path, "a") as f:
        f.write(string + "\n")

runID = str(int(round(time.time(),0)))

cluster_kwargs = {
    "name": "exDaskCluster",
    "shebang": "#!/bin/tcsh",
    "resource_spec": "nodes=1:c18a:ppn=12",
    "walltime": "01:00:00",
    "cores": 12,
    "processes": 12,
    "memory": "32GB",
    "interface": "ib0",
}


cluster = PBSCluster(**cluster_kwargs)
cluster.scale(PROCESSES)
client = Client(cluster)
client.wait_for_workers(PROCESSES)

rank = [client.scatter(rank) for rank in range(2)]
world_size = [client.scatter(2) for _ in range(2)]
logger(rank)
logger(world_size)


#Network
class NeuralNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten()
        self.simpleNet = nn.Sequential(
            nn.Linear(64*64*3, 512),
            nn.ReLU(),
            nn.Linear(512, 128),
            nn.ReLU(),
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Linear(64, 32),
            nn.Linear(32, 21)
        )

    def forward(self, x):
        x = self.flatten(x)
        x = self.simpleNet(x)
        return x

def initDist(rank, size, runID, EPOCHS_COUNT, backend='gloo'):
    from torchvision import transforms

    BATCH_SIZE = 4

    """
    Initialize the distributed environment
    """
    if(rank == 0):
        ip = socket.gethostbyname(socket.gethostname())
        with open(BASE_PATH+"ip_"+ runID +".dan", "w") as f:
            f.write(ip)
    else:
        ip = None
        while ip == None:
            try:
                with open(BASE_PATH+"ip_"+ runID + ".dan", "r") as f:
                    ip = f.readlines()[0]
            except:
                print("Waiting for master node to connect.")
                time.sleep(1)

    def logger(string, path=BASE_PATH+"log.dan"):
        string = str(string)
        with open(path, "a") as f:
            f.write(string + "\n")

    os.environ['MASTER_ADDR'] = ip #if using multiple workers, set the address of the worker with rank = 0
    os.environ['MASTER_PORT'] = "23456"
    os.environ['WORLD_SIZE'] = str(size)
    os.environ['RANK'] = str(rank)

    ## Init distributed
    logger(str(rank) + ": Waiting for all workers to connect...")
    dist.init_process_group(
        init_method="env://",
        backend='gloo'
    )
    device = torch.device('cpu')  
    dist.barrier()

    #At this point, our cluster is initialized - i.e., 
    #the dist.init_process_group will have run on all nodes.
    #Code is blocked from proceeding until that happens.
    #We can now use special torch functions that only work
    #in distributed cases.

    #First, we need to distribute our actual model architecture to all nodes.
    #This is only done once.
    logger(str(rank) + ": Loading model.")
    model = models.resnet101(pretrained=False)
    ddp_model = DDP(model, device_ids=None) #For GPUs, you have to specify device IDs here.
    logger(str(rank) + ": Model loaded.")

    #Now, we need to distribute our data across all the nodes.
    #This is done by creating sub-samples and sending a sub-sample to each node.
    #DistributedSampler handles this process.
    logger(str(rank) + ": Loading imagery.")
    transforms = transforms.Compose([transforms.Resize((64,64)),
                                 transforms.ToTensor()])
    images = datasets.ImageFolder(IMAGE_PATH, transform=transforms)
    s = DistributedSampler(images)
    loader = DataLoader(images, batch_size=BATCH_SIZE, sampler=s)
    logger(str(rank) + " Dataset size:" + str(len(loader)))

    #Now we can (finally) train our model like usual!
    lossFN = nn.CrossEntropyLoss()
    optimizer = optim.SGD(ddp_model.parameters(), lr=.01)
    epochs = EPOCHS_COUNT
    num_of_batches = len(loader)
    
    for epoch in range(epochs):
        logger(str(rank) + ": starting epoch " + str(epoch) + " | Batches in Epoch: " + str(num_of_batches))

        #This allows us to shuffle our data
        loader.sampler.set_epoch(epoch)

        ddp_model.train()

        for batch, (X,y) in enumerate(loader):
            X, y = X.to(device), y.to(device)
            #Make our predictions and calculate loss
            forwardPassPred = ddp_model(X)
            loss = lossFN(forwardPassPred, y)

            #Backpropogate
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            if batch % 10 == 0:
                loss, current = loss.item(), batch*len(X)
                logger(str(rank) + "||  Loss: " + str(loss) + " | " + str(current) + " of " + str(len(X) * num_of_batches))

        if rank == 0:
            torch.save(ddp_model.module.state_dict(),BASE_PATH+"checkpoints/" + str(runID) + "_" + str(epoch)+".checkpoint")

        #Prevents any node from continuing to optimize our model before it's done saving as the end of an epoch.
        dist.barrier()

    dist.destroy_process_group()
    #You can return anything you want here - we've saved our weights already.
    #This just goes back to Dask to confirm the process completed.
    #If you added "print(results)" after the client.gather, you see these returns.
    return(str(loss))

future = client.map(initDist, list(range(PROCESSES)), [PROCESSES]*PROCESSES, [runID]*PROCESSES, [EPOCHS_COUNT]*PROCESSES)
logger(future)
results = client.gather(future)

#And, we're done with Dask!  We can now explore our model fit itself, but shut down our
#Dask nodes as they aren't doing anything now:
client.close()

#Now that we have our model, we can load it and run accuracy statistics on it.
#We're just going to run this on our "master" node - i.e., the same node we used
#to spinup the Dask pieces.
from torchvision import transforms
import pandas as pd

def accuracyStatistics(model, dataLoader):
    model.eval()
    test_loss = 0
    device = 'cpu'

    y_pred = []
    y_actu = []
    with torch.no_grad():
        for data, target in dataLoader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            test_loss += F.nll_loss(output, target, reduction='sum').item()  # sum up batch loss
            pred = output.argmax(dim=1, keepdim=True)  # get the index of the max score
            y_pred.extend(torch.flatten(pred).tolist()) 
            y_actu.extend(target.tolist())
           
    y_pred = pd.Series(y_pred, name='Predicted')
    y_actu = pd.Series(y_actu, name='Actual')
    cm = pd.crosstab(y_actu, y_pred)

    return(cm)

#Note we're not using the DDP model or loader here,
#as everything is now just running (in forward-pass mode) 
#on the single node.  Also no sampler for our dataset, 
#as we want to test agains the whole thing.
#(Note: in practice we would want a test/train split,
#but I haven't created two datasets for that purpose here.)
model = models.resnet101(pretrained=False)

#Now we need to load the weights for the network we just solved in Dask-world:
model.load_state_dict(torch.load(BASE_PATH+"checkpoints/" + str(runID) + "_" + str(EPOCHS_COUNT-1)+".checkpoint"))

#Load our image data:
transforms = transforms.Compose([transforms.Resize((64,64)),
                                 transforms.ToTensor()])
images = datasets.ImageFolder(IMAGE_PATH, transform=transforms)
loader = DataLoader(images, batch_size=4)

#Create our crosstabulation matrix:
logger(accuracyStatistics(model, loader))





print("======SHUTDOWN INITIALIZING=======")
print("==================================")
print("==================================")
print("==================================")
print("==================================")
print("==================================")
print("==================================")
print("==================================")
print("==================================")
print("==================================")
print("==================================")
print("==================================")
print("==================================")
print("======SHUTDOWN INITIALIZING=======")