xu-ji

# run before you do anything else in the code, for each experiment

def set_seed(seed):
  random.seed(0)
  np.random.seed(0)
  torch.manual_seed(seed)
  torch.cuda.manual_seed(seed)
  torch.backends.cudnn.deterministic = True
  torch.backends.cudnn.benchmark = False

xu-ji

from pynvml.smi import nvidia_smi

def get_gpu_mem(nvsmi, gpu_ind):
  mem_stats = nvsmi.DeviceQuery('memory.free, memory.total')["gpu"][gpu_ind]["fb_memory_usage"]
  return mem_stats["total"] - mem_stats["free"]

nvsmi = nvidia_smi.getInstance()
gpu_ind = 0
print(get_gpu_mem(nvsmi, gpu_ind))

xu-ji

from code.utils.cluster.cluster_eval import cluster_eval
from code.utils.cluster.data import cluster_twohead_create_dataloaders
import pickle
import code.archs as archs
import torch
from datetime import datetime
import sys

config_in = open("/scratch/shared/nfs1/xuji/iid_private/685/config.pickle", "rb")
config = pickle.load(config_in)