rushi-the-neural-arch

def save_latest(epoch, model_dir, model, optimizer, scheduler=None, wandb_id=None):
        """Saves latest epoch's weights and other necessary things to resume
        """
        model_states = {
            "epoch": epoch,
            "state_dict": model.state_dict(),
            "opt_state_dict": optimizer.state_dict(),
            "sch_state_dict": scheduler.state_dict() if scheduler != None else None,
            "wandb_id_save": wandb_id,        #----> Remove this if you don't use wandb for logging
        }

rushi-the-neural-arch

import torch
import torch.nn as nn
import torch.distributed as dist
from torch.utils.data import DataLoader
import torch.multiprocessing as mp
from fairscale.optim.oss import OSS
from fairscale.nn.data_parallel import ShardedDataParallel as ShardedDDP

from torch.optim import AdamW

rushi-the-neural-arch

#python -m torch.distributed.launch Stoke-DDP.py --projectName "PyTorch-4K-2X" --batchSize 20 --nEpochs 2 --lr 1e-3 --threads 8
#env CUDA_VISIBLE_DEVICES=0,1,2,3 python -m torch.distributed.launch --nproc_per_node=4 Stoke-DDP.py --projectName "Stoke-4K-2X-DDP" --batchSize 18 --nEpochs 2 --lr 1e-3 --weight_decay 1e-4 --grad_clip 0.1

import argparse, os, sys
import random
import numpy as np
import time
import torch
import torch.nn as nn
import torch.backends.cudnn as cudnn 

rushi-the-neural-arch

# python3 DDP.py --projectName "PyTorch-4K-2X" --batchSize 32 --nEpochs 2 --lr 1e-3 --step 10 --threads 8 --optim AdamW

import argparse, os, sys
import random
import numpy as np
import time
import torch
import torch.nn as nn
import torch.backends.cudnn as cudnn 
import torch.optim as optim

rushi-the-neural-arch

import os
import heapq
import collections
import operator
import ast
import sys
import time


class HeapNode: