def setup(rank, world_size):
    # initialize the process group
    dist.init_process_group("nccl", rank=rank, world_size=world_size)

model = Net().to(device)
optimizer = optim.Adadelta(model.parameters(), lr=args.lr)

model = Net().to(local_rank)
ddp_model = DDP(model, device_ids=[local_rank])
optimizer = optim.Adadelta(ddp_model.parameters(), lr=args.lr)

train_sampler = torch.utils.data.distributed.DistributedSampler(\
                            dataset1, num_replicas=world_size, rank=rank)
train_loader  = torch.utils.data.DataLoader(\
                            dataset1, batch_size=args.batch_size,  \
                            sampler=train_sampler, \ 
                            num_workers=int(os.environ["SLURM_CPUS_PER_TASK"]), \
                            pin_memory=True)

#SBATCH -p dgx-preempt           # trainning partition ที่เราเลือกใช้คือ dgx-preempt
#SBATCH -N 1                     # จำนวน node ของ dgx-preempt ที่เราต้องการ
#SBATCH --ntasks-per-node=8      # จำนวน ntasks per node (=จำนวน gpus ที่มี)
#SBATCH --cpus-per-task=5        # จำนวน cpu per task (5x8=40 cpus)
...
export WORLD_SIZE=8   # ควรได้มาจาก $(($SLURM_NNODES * $SLURM_NTASKS_PER_NODE))
echo "WORLD_SIZE="$WORLD_SIZE
...
srun python DDP.py

    JOBID PARTITION     NAME     USER ST       TIME  NODES NODELIST(REASON)
   0000001       dgx  myjob   xxxxxxx PD       0:00      1 (PartitionConfig)
   0000002   compute herjob   yyyyyyy PD       0:00      1 (QOSGrpBillingMinutes)
   0000003 dgx-preem hisjob   xxxxxxx PD       0:00      1 (Resources)
   0000004       gpu onejob   yyyyyyy PD       0:00      1 (Priority)
   0000005       gpu onejob   yyyyyyy  R    6:40:53      1 tara-g-001

import os
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from socket import gethostname

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc = nn.Linear(42, 3)

    def forward(self, x):
        x = self.fc(x)
        x = F.relu(x)
        output = F.log_softmax(x, dim=1)
        return output

rank          = int(os.environ["SLURM_PROCID"])
world_size    = int(os.environ["WORLD_SIZE"])
gpus_per_node = int(os.environ["SLURM_GPUS_ON_NODE"])
assert gpus_per_node == torch.cuda.device_count()
print(f"Hello from rank {rank} of {world_size} on {gethostname()} where there are" \
      f" {gpus_per_node} allocated GPUs per node.", flush=True)

dist.init_process_group("nccl", rank=rank, world_size=world_size)
if rank == 0: print(f"Group initialized? {dist.is_initialized()}", flush=True)

local_rank = rank - gpus_per_node * (rank // gpus_per_node)
torch.cuda.set_device(local_rank)

model = Net().to(local_rank)
ddp_model = DDP(model, device_ids=[local_rank])

ddp_model.eval()
with torch.no_grad():
  data = torch.rand(1, 42)
  data = data.to(local_rank)
  output = ddp_model(data)
  print(f"host: {gethostname()}, rank: {rank}, output: {output}")

dist.destroy_process_group()

rank          = int(os.environ["SLURM_PROCID"])
world_size    = int(os.environ["WORLD_SIZE"])
gpus_per_node = int(os.environ["SLURM_GPUS_ON_NODE"])
...
local_rank = rank - gpus_per_node * (rank // gpus_per_node)
torch.cuda.set_device(local_rank)

#SBATCH -N 2
#SBATCH --ntasks-per-node=4
...
srun python myscript.py

WORLD_SIZE=8
MASTER_PORT=15334
MASTER_ADDR=tara-dgx1-002
Hello from rank 6 of 8 on tara-dgx1-002.tara.nstda.or.th where there are 8 allocated GPUs per node.
Hello from rank 5 of 8 on tara-dgx1-002.tara.nstda.or.th where there are 8 allocated GPUs per node.
Hello from rank 4 of 8 on tara-dgx1-002.tara.nstda.or.th where there are 8 allocated GPUs per node.
Hello from rank 7 of 8 on tara-dgx1-002.tara.nstda.or.th where there are 8 allocated GPUs per node.
Hello from rank 3 of 8 on tara-dgx1-002.tara.nstda.or.th where there are 8 allocated GPUs per node.
Hello from rank 0 of 8 on tara-dgx1-002.tara.nstda.or.th where there are 8 allocated GPUs per node.
Hello from rank 1 of 8 on tara-dgx1-002.tara.nstda.or.th where there are 8 allocated GPUs per node.
Hello from rank 2 of 8 on tara-dgx1-002.tara.nstda.or.th where there are 8 allocated GPUs per node.
Group initialized? True
host: tara-dgx1-002.tara.nstda.or.th, rank: 6, output: tensor([[-1.2452, -1.0438, -1.0216]], device='cuda:6')
host: tara-dgx1-002.tara.nstda.or.th, rank: 4, output: tensor([[-1.2201, -0.9631, -1.1298]], device='cuda:4')
host: tara-dgx1-002.tara.nstda.or.th, rank: 5, output: tensor([[-1.2976, -0.8704, -1.1775]], device='cuda:5')
host: tara-dgx1-002.tara.nstda.or.th, rank: 2, output: tensor([[-1.2821, -1.0588, -0.9790]], device='cuda:2')
host: tara-dgx1-002.tara.nstda.or.th, rank: 1, output: tensor([[-1.1647, -0.9979, -1.1416]], device='cuda:1')
host: tara-dgx1-002.tara.nstda.or.th, rank: 0, output: tensor([[-1.1372, -1.0257, -1.1372]], device='cuda:0')
host: tara-dgx1-002.tara.nstda.or.th, rank: 3, output: tensor([[-1.2476, -1.0863, -0.9799]], device='cuda:3')
host: tara-dgx1-002.tara.nstda.or.th, rank: 7, output: tensor([[-1.1277, -1.0428, -1.1277]], device='cuda:7')

#!/bin/bash
#SBATCH -A thaisc                # account of your project
#SBATCH -J SG-torch              # create a short name for your job
#SBATCH -p dgx-preempt           # your choice of partition
#SBATCH -N 1                     # node count
#SBATCH --ntasks-per-node=8      # total number of tasks per node
#SBATCH --time=00:05:00          # total run time limit (HH:MM:SS)
#SBATCH --mail-type=ALL          # send email when job begins, ends
#SBATCH --mail-user=me@myorg.org

export MASTER_PORT=$(expr 10000 + $(echo -n $SLURM_JOBID | tail -c 4))
echo "MASTER_PORT="$MASTER_PORT

export WORLD_SIZE=8   # ควรได้มาจาก $(($SLURM_NNODES * $SLURM_NTASKS_PER_NODE))
echo "WORLD_SIZE="$WORLD_SIZE

master_addr=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
export MASTER_ADDR=$master_addr
echo "MASTER_ADDR="$MASTER_ADDR

module purge
module load cuDNN/8.0.5.39-CUDA-11.1.1
source ~/.bashrc
myconda
conda activate condapy37

srun python simpleDDP.py

import argparse
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
from torch.optim.lr_scheduler import StepLR

import os
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from socket import gethostname

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.dropout1 = nn.Dropout(0.25)
        self.dropout2 = nn.Dropout(0.5)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2)
        x = self.dropout1(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.dropout2(x)
        x = self.fc2(x)
        output = F.log_softmax(x, dim=1)
        return output

def train(args, model, device, train_loader, optimizer, epoch):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = F.nll_loss(output, target)
        loss.backward()
        optimizer.step()
        if batch_idx % args.log_interval == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                100. * batch_idx / len(train_loader), loss.item()))
            if args.dry_run:
                break

def test(model, device, test_loader):
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            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 log-probability
            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(test_loader.dataset)

    print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
        test_loss, correct, len(test_loader.dataset),
        100. * correct / len(test_loader.dataset)))

def setup(rank, world_size):
    # initialize the process group
    dist.init_process_group("nccl", rank=rank, world_size=world_size)

def main():
    # Training settings
    parser = argparse.ArgumentParser(description='PyTorch MNIST Example')
    parser.add_argument('--batch-size', type=int, default=64, metavar='N',
                        help='input batch size for training (default: 64)')
    parser.add_argument('--test-batch-size', type=int, default=1000, metavar='N',
                        help='input batch size for testing (default: 1000)')
    parser.add_argument('--epochs', type=int, default=14, metavar='N',
                        help='number of epochs to train (default: 14)')
    parser.add_argument('--lr', type=float, default=1.0, metavar='LR',
                        help='learning rate (default: 1.0)')
    parser.add_argument('--gamma', type=float, default=0.7, metavar='M',
                        help='Learning rate step gamma (default: 0.7)')
    parser.add_argument('--no-cuda', action='store_true', default=False,
                        help='disables CUDA training')
    parser.add_argument('--dry-run', action='store_true', default=False,
                        help='quickly check a single pass')
    parser.add_argument('--seed', type=int, default=1, metavar='S',
                        help='random seed (default: 1)')
    parser.add_argument('--log-interval', type=int, default=10, metavar='N',
                        help='how many batches to wait before logging training status')
    parser.add_argument('--save-model', action='store_true', default=False,
                        help='For Saving the current Model')
    args = parser.parse_args()
    use_cuda = not args.no_cuda and torch.cuda.is_available()

    torch.manual_seed(args.seed)
    
    train_kwargs = {'batch_size': args.batch_size}
    test_kwargs = {'batch_size': args.test_batch_size}
    if use_cuda:
        cuda_kwargs = {'num_workers': int(os.environ["SLURM_CPUS_PER_TASK"]),
                       'pin_memory': True,
                       'shuffle': True}
        train_kwargs.update(cuda_kwargs)
        test_kwargs.update(cuda_kwargs)

    transform=transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
        ])
    dataset1 = datasets.MNIST('data', train=True, download=False,
                       transform=transform)
    dataset2 = datasets.MNIST('data', train=False,
                       transform=transform)
 
    world_size    = int(os.environ["WORLD_SIZE"])
    rank          = int(os.environ["SLURM_PROCID"])
    gpus_per_node = int(os.environ["SLURM_GPUS_ON_NODE"])
    assert gpus_per_node == torch.cuda.device_count()
    print(f"Hello from rank {rank} of {world_size} on {gethostname()} where there are" \
          f" {gpus_per_node} allocated GPUs per node.", flush=True)

    setup(rank, world_size)
    if rank == 0: print(f"Group initialized? {dist.is_initialized()}", flush=True)

    local_rank = rank - gpus_per_node * (rank // gpus_per_node)
    torch.cuda.set_device(local_rank)
    print(f"host: {gethostname()}, rank: {rank}, local_rank: {local_rank}")

    train_sampler = torch.utils.data.distributed.DistributedSampler(dataset1, num_replicas=world_size, rank=rank)
    train_loader = torch.utils.data.DataLoader(dataset1, batch_size=args.batch_size, sampler=train_sampler, \
                                               num_workers=int(os.environ["SLURM_CPUS_PER_TASK"]), pin_memory=True)
    test_loader = torch.utils.data.DataLoader(dataset2, **test_kwargs)

    model = Net().to(local_rank)
    ddp_model = DDP(model, device_ids=[local_rank])
    optimizer = optim.Adadelta(ddp_model.parameters(), lr=args.lr)

    scheduler = StepLR(optimizer, step_size=1, gamma=args.gamma)
    for epoch in range(1, args.epochs + 1):
        train(args, ddp_model, local_rank, train_loader, optimizer, epoch)
        if rank == 0: test(ddp_model, local_rank, test_loader)
        scheduler.step()

    if args.save_model and rank == 0:
        torch.save(model.state_dict(), "mnist_cnn.pt")

    dist.destroy_process_group()


if __name__ == '__main__':
    main()

#!/bin/bash
#SBATCH -A thaisc                # account of your project
#SBATCH -J SG-torch              # create a short name for your job
#SBATCH -p dgx-preempt           # your choice of partition
#SBATCH -N 1                     # node count
#SBATCH --ntasks-per-node=8      # total number of tasks per node
#SBATCH --cpus-per-task=5        # number of cpu per task (5x8=40 cpus)
#SBATCH --time=00:05:00          # total run time limit (HH:MM:SS)
#SBATCH --mail-type=ALL          # send email when job begins, ends
#SBATCH --mail-user=me@myorg.org

export MASTER_PORT=$(expr 10000 + $(echo -n $SLURM_JOBID | tail -c 4))
echo "MASTER_PORT="$MASTER_PORT

export WORLD_SIZE=8    # ควรได้มาจาก $(($SLURM_NNODES * $SLURM_NTASKS_PER_NODE))
echo "WORLD_SIZE="$WORLD_SIZE

master_addr=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
export MASTER_ADDR=$master_addr
echo "MASTER_ADDR="$MASTER_ADDR

module purge
module load cuDNN/8.0.5.39-CUDA-11.1.1
source ~/.bashrc
myconda
conda activate condapy37

srun python DDP.py

WORLD_SIZE=8
MASTER_PORT=15347
MASTER_ADDR=tara-dgx1-002
Hello from rank 1 of 8 on tara-dgx1-002.tara.nstda.or.th where there are 8 allocated GPUs per node.
Hello from rank 3 of 8 on tara-dgx1-002.tara.nstda.or.th where there are 8 allocated GPUs per node.
Hello from rank 0 of 8 on tara-dgx1-002.tara.nstda.or.th where there are 8 allocated GPUs per node.
Hello from rank 2 of 8 on tara-dgx1-002.tara.nstda.or.th where there are 8 allocated GPUs per node.
Hello from rank 5 of 8 on tara-dgx1-002.tara.nstda.or.th where there are 8 allocated GPUs per node.
Hello from rank 4 of 8 on tara-dgx1-002.tara.nstda.or.th where there are 8 allocated GPUs per node.
Hello from rank 6 of 8 on tara-dgx1-002.tara.nstda.or.th where there are 8 allocated GPUs per node.
Hello from rank 7 of 8 on tara-dgx1-002.tara.nstda.or.th where there are 8 allocated GPUs per node.
Group initialized? True
host: tara-dgx1-002.tara.nstda.or.th, rank: 1, local_rank: 1
host: tara-dgx1-002.tara.nstda.or.th, rank: 3, local_rank: 3
host: tara-dgx1-002.tara.nstda.or.th, rank: 4, local_rank: 4
host: tara-dgx1-002.tara.nstda.or.th, rank: 6, local_rank: 6
host: tara-dgx1-002.tara.nstda.or.th, rank: 0, local_rank: 0
host: tara-dgx1-002.tara.nstda.or.th, rank: 2, local_rank: 2
host: tara-dgx1-002.tara.nstda.or.th, rank: 5, local_rank: 5
host: tara-dgx1-002.tara.nstda.or.th, rank: 7, local_rank: 7
Train Epoch: 1 [0/60000 (0%)]   Loss: 2.297117
Train Epoch: 1 [640/60000 (8%)] Loss: 1.329343
Train Epoch: 1 [1280/60000 (17%)]       Loss: 0.518520
Train Epoch: 1 [1920/60000 (25%)]       Loss: 0.331641
Train Epoch: 1 [2560/60000 (34%)]       Loss: 0.256029
Train Epoch: 1 [3200/60000 (42%)]       Loss: 0.126544
Train Epoch: 1 [3840/60000 (51%)]       Loss: 0.129393
Train Epoch: 1 [4480/60000 (59%)]       Loss: 0.135831
Train Epoch: 1 [5120/60000 (68%)]       Loss: 0.094554
Train Epoch: 1 [5760/60000 (76%)]       Loss: 0.131771
Train Epoch: 1 [6400/60000 (85%)]       Loss: 0.078105
Train Epoch: 1 [7040/60000 (93%)]       Loss: 0.078772
Train Epoch: 1 [0/60000 (0%)]   Loss: 2.315368
Train Epoch: 1 [640/60000 (8%)] Loss: 1.471632
Train Epoch: 1 [1280/60000 (17%)]       Loss: 0.394169
Train Epoch: 1 [1920/60000 (25%)]       Loss: 0.376319
...
Train Epoch: 14 [5120/60000 (68%)]      Loss: 0.003920
Train Epoch: 14 [5760/60000 (76%)]      Loss: 0.105166
Train Epoch: 14 [6400/60000 (85%)]      Loss: 0.020963
Train Epoch: 14 [7040/60000 (93%)]      Loss: 0.071237

Test set: Average loss: 0.0298, Accuracy: 9897/10000 (99%)

#SBATCH -N 2                     # node count
...
export WORLD_SIZE=16    # ควรได้มาจาก $(($SLURM_NNODES * $SLURM_NTASKS_PER_NODE))

WORLD_SIZE=16
MASTER_PORT=17117
MASTER_ADDR=tara-dgx1-002
Hello from rank 12 of 16 on tara-dgx1-003.tara.nstda.or.th where there are 8 allocated GPUs per node.
Hello from rank 13 of 16 on tara-dgx1-003.tara.nstda.or.th where there are 8 allocated GPUs per node.
Hello from rank 15 of 16 on tara-dgx1-003.tara.nstda.or.th where there are 8 allocated GPUs per node.
Hello from rank 8 of 16 on tara-dgx1-003.tara.nstda.or.th where there are 8 allocated GPUs per node.
Hello from rank 9 of 16 on tara-dgx1-003.tara.nstda.or.th where there are 8 allocated GPUs per node.
Hello from rank 10 of 16 on tara-dgx1-003.tara.nstda.or.th where there are 8 allocated GPUs per node.
Hello from rank 11 of 16 on tara-dgx1-003.tara.nstda.or.th where there are 8 allocated GPUs per node.
Hello from rank 14 of 16 on tara-dgx1-003.tara.nstda.or.th where there are 8 allocated GPUs per node.
Hello from rank 4 of 16 on tara-dgx1-002.tara.nstda.or.th where there are 8 allocated GPUs per node.
Hello from rank 5 of 16 on tara-dgx1-002.tara.nstda.or.th where there are 8 allocated GPUs per node.
Hello from rank 6 of 16 on tara-dgx1-002.tara.nstda.or.th where there are 8 allocated GPUs per node.
Hello from rank 7 of 16 on tara-dgx1-002.tara.nstda.or.th where there are 8 allocated GPUs per node.
Hello from rank 0 of 16 on tara-dgx1-002.tara.nstda.or.th where there are 8 allocated GPUs per node.
Hello from rank 1 of 16 on tara-dgx1-002.tara.nstda.or.th where there are 8 allocated GPUs per node.
Hello from rank 2 of 16 on tara-dgx1-002.tara.nstda.or.th where there are 8 allocated GPUs per node.
Hello from rank 3 of 16 on tara-dgx1-002.tara.nstda.or.th where there are 8 allocated GPUs per node.
Group initialized? True
host: tara-dgx1-002.tara.nstda.or.th, rank: 4, local_rank: 4
host: tara-dgx1-002.tara.nstda.or.th, rank: 6, local_rank: 6
host: tara-dgx1-002.tara.nstda.or.th, rank: 0, local_rank: 0
host: tara-dgx1-002.tara.nstda.or.th, rank: 2, local_rank: 2
host: tara-dgx1-002.tara.nstda.or.th, rank: 5, local_rank: 5
host: tara-dgx1-002.tara.nstda.or.th, rank: 7, local_rank: 7
host: tara-dgx1-002.tara.nstda.or.th, rank: 1, local_rank: 1
host: tara-dgx1-002.tara.nstda.or.th, rank: 3, local_rank: 3
host: tara-dgx1-003.tara.nstda.or.th, rank: 12, local_rank: 4
host: tara-dgx1-003.tara.nstda.or.th, rank: 13, local_rank: 5
host: tara-dgx1-003.tara.nstda.or.th, rank: 15, local_rank: 7
host: tara-dgx1-003.tara.nstda.or.th, rank: 8, local_rank: 0
host: tara-dgx1-003.tara.nstda.or.th, rank: 9, local_rank: 1
host: tara-dgx1-003.tara.nstda.or.th, rank: 10, local_rank: 2
host: tara-dgx1-003.tara.nstda.or.th, rank: 11, local_rank: 3
host: tara-dgx1-003.tara.nstda.or.th, rank: 14, local_rank: 6
Train Epoch: 1 [0/60000 (0%)]   Loss: 2.288294
Train Epoch: 1 [640/60000 (17%)]        Loss: 1.147667
Train Epoch: 1 [1280/60000 (34%)]       Loss: 0.518015
Train Epoch: 1 [1920/60000 (51%)]       Loss: 0.241061
Train Epoch: 1 [2560/60000 (68%)]       Loss: 0.159519
Train Epoch: 1 [3200/60000 (85%)]       Loss: 0.387764
Train Epoch: 1 [0/60000 (0%)]   Loss: 2.306850
Train Epoch: 1 [640/60000 (17%)]        Loss: 1.047796
...

$ conda activate condapy37
(condapy37)$ python
Python 3.7.13 (default, Mar 29 2022, 02:18:16)
[GCC 7.5.0] :: Anaconda, Inc. on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> from torchvision import datasets, transforms
>>> transform=transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
        ])
>>> dataset1 = datasets.MNIST('data', train=True, download=True, transform=transform)

Downloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz
Downloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz to data/MNIST/raw/train-images-idx3-ubyte.gz
9913344it [00:01, 5744180.56it/s]
Extracting data/MNIST/raw/train-images-idx3-ubyte.gz to data/MNIST/raw

Downloading http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz
Downloading http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz to data/MNIST/raw/train-labels-idx1-ubyte.gz
29696it [00:00, 1196864.05it/s]
Extracting data/MNIST/raw/train-labels-idx1-ubyte.gz to data/MNIST/raw

Downloading http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz
Downloading http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz to data/MNIST/raw/t10k-images-idx3-ubyte.gz
1649664it [00:00, 2396549.23it/s]
Extracting data/MNIST/raw/t10k-images-idx3-ubyte.gz to data/MNIST/raw

Downloading http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz
Downloading http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz to data/MNIST/raw/t10k-labels-idx1-ubyte.gz
5120it [00:00, 23165950.90it/s]
Extracting data/MNIST/raw/t10k-labels-idx1-ubyte.gz to data/MNIST/raw

>>> exit()