def safeunpickle(path):
if path.startswith("jittorhub://"):
path = path.replace(
"jittorhub://",
"https://cg.cs.tsinghua.edu.cn/jittor/assets/build/checkpoints/")
if path.startswith("https:") or path.startswith("http:"):
base = path.split("/")[-1]
fname = os.path.join(compiler.ck_path, base)
from jittor_utils.misc import download_url_to_local
download_url_to_local(path, base, compiler.ck_path, None)
path = fname
if path.endswith(".pth"):
try:
dirty_fix_pytorch_runtime_error()
import torch
except:
raise RuntimeError(
"pytorch need to be installed when load pth format.")
model_dict = torch.load(path, map_location=torch.device('cpu'))
return model_dict
with open(path, "rb") as f:
s = f.read()
if not s.endswith(b"HCAJSLHD"):
return pickle.loads(s)
checksum = s[-28:-8]
s = s[:-28]
if hashlib.sha1(s).digest() != checksum:
raise ValueError("Pickle checksum does not match! path: " + path)
return pickle.loads(s)
def install(path):
LOG.i("Installing MSVC...")
filename = "msvc.zip"
url = "https://cg.cs.tsinghua.edu.cn/jittor/assets/" + filename
md5sum = "55f0c175fdf1419b124e0fc498b659d2"
download_url_to_local(url, filename, path, md5sum)
fullname = os.path.join(path, filename)
import zipfile
with zipfile.ZipFile(fullname, "r") as f:
f.extractall(path)
def install_mkl(root_folder):
# origin url is
# url = "https://github.com/intel/mkl-dnn/releases/download/v1.0.2/mkldnn_lnx_1.0.2_cpu_gomp.tgz"
import platform
url = None
if platform.system()=="Linux":
if platform.machine()=='x86_64':
filename = "dnnl_lnx_2.2.0_cpu_gomp.tgz"
md5 = "35bbbdf550a9d8ad54db798e372000f6"
elif platform.machine()=='aarch64':
filename = "dnnl_lnx_2.2.0_cpu_gomp_aarch64.tgz"
md5 = "72cf9b0b8fd6c3c786d35a9daaee22b8"
else:
raise RuntimeError(f"platform.machine()=={platform.machine()} not support yet,"
" Please contact us on https://github.com/jittor/jittor ")
elif os.name == "nt":
# url = "https://github.com/oneapi-src/oneDNN/releases/download/v2.2/dnnl_win_2.2.0_cpu_iomp.zip"
# url = "https://github.com/oneapi-src/oneDNN/releases/download/v2.2/dnnl_win_2.2.0_cpu_vcomp.zip"
filename = "dnnl_win_2.2.0_cpu_vcomp.zip"
md5 = "fa12c693b2ec07700d174e1e99d60a7e"
else:
raise RuntimeError(f"platform.machine()=={platform.machine()} not support yet,"
" Please contact us on https://github.com/jittor/jittor ")
if not url:
url = "https://cg.cs.tsinghua.edu.cn/jittor/assets/" + filename
fullname = os.path.join(root_folder, filename)
dirname = os.path.join(root_folder, filename.rsplit(".",1)[0])
if not (os.path.isfile(os.path.join(dirname, "lib", "libmkldnn.so")) or
os.path.isfile(os.path.join(dirname, "bin", "dnnl.dll"))):
LOG.i("Downloading mkl...")
download_url_to_local(url, filename, root_folder, md5)
if fullname.endswith(".zip"):
import zipfile
with zipfile.ZipFile(fullname, "r") as f:
f.extractall(root_folder)
else:
import tarfile
with tarfile.open(fullname, "r") as tar:
tar.extractall(root_folder)
if os.name == 'nt':
# this env is used for execute example/text
bin_path = os.path.join(dirname, "bin")
sys.path.append(bin_path)
os.environ["PATH"] = os.environ.get("PATH", "") + ";" + bin_path
cmd = f"cd /d {dirname}/examples && {cc_path} {dirname}/examples/cnn_inference_f32.cpp -I{dirname}/include -Fe: {dirname}/examples/test.exe {fix_cl_flags(cc_flags).replace('-LD', '')} {dirname}/lib/mkldnn.lib"
assert 0 == os.system(cmd)
assert 0 == os.system(f"{dirname}/examples/test")
else:
assert 0 == os.system(f"cd {dirname}/examples && "
f"{cc_path} -std=c++14 cnn_inference_f32.cpp -Ofast -lmkldnn -I ../include -L ../lib -o test && LD_LIBRARY_PATH=../lib/ ./test")
def download_url(self):
'''
Download mnist data set function, this function will be called when download is True.
'''
resources = [
("https://storage.googleapis.com/cvdf-datasets/mnist/train-images-idx3-ubyte.gz", "f68b3c2dcbeaaa9fbdd348bbdeb94873"),
("https://storage.googleapis.com/cvdf-datasets/mnist/train-labels-idx1-ubyte.gz", "d53e105ee54ea40749a09fcbcd1e9432"),
("https://storage.googleapis.com/cvdf-datasets/mnist/t10k-images-idx3-ubyte.gz", "9fb629c4189551a2d022fa330f9573f3"),
("https://storage.googleapis.com/cvdf-datasets/mnist/t10k-labels-idx1-ubyte.gz", "ec29112dd5afa0611ce80d1b7f02629c")
]
for url, md5 in resources:
filename = url.rpartition('/')[2]
download_url_to_local(url, filename, self.data_root, md5)
def install_cuda():
cuda_driver_version = get_cuda_driver()
if not cuda_driver_version:
return None
LOG.i("cuda_driver_version: ", cuda_driver_version)
if cuda_driver_version >= [11, 2]:
cuda_tgz = "cuda11.2_cudnn8_linux.tgz"
md5 = "b93a1a5d19098e93450ee080509e9836"
elif cuda_driver_version >= [
11,
]:
cuda_tgz = "cuda11.0_cudnn8_linux.tgz"
md5 = "5dbdb43e35b4db8249027997720bf1ca"
elif cuda_driver_version >= [10, 2]:
cuda_tgz = "cuda10.2_cudnn7_linux.tgz"
md5 = "40f0563e8eb176f53e55943f6d212ad7"
elif cuda_driver_version >= [
10,
]:
cuda_tgz = "cuda10.0_cudnn7_linux.tgz"
md5 = "f16d3ff63f081031d21faec3ec8b7dac"
else:
raise RuntimeError(
f"Unsupport cuda driver version: {cuda_driver_version}")
jtcuda_path = os.path.join(pathlib.Path.home(), ".cache", "jittor",
"jtcuda")
nvcc_path = os.path.join(jtcuda_path, cuda_tgz[:-4], "bin", "nvcc")
nvcc_lib_path = os.path.join(jtcuda_path, cuda_tgz[:-4], "lib64")
sys.path.append(nvcc_lib_path)
new_ld_path = os.environ.get("LD_LIBRARY_PATH", "") + ":" + nvcc_lib_path
os.environ["LD_LIBRARY_PATH"] = new_ld_path
if os.path.isfile(nvcc_path):
return nvcc_path
os.makedirs(jtcuda_path, exist_ok=True)
cuda_tgz_path = os.path.join(jtcuda_path, cuda_tgz)
download_url_to_local(
"https://cg.cs.tsinghua.edu.cn/jittor/assets/" + cuda_tgz, cuda_tgz,
jtcuda_path, md5)
import tarfile
with tarfile.open(cuda_tgz_path, "r") as tar:
tar.extractall(cuda_tgz_path[:-4])
assert os.path.isfile(nvcc_path)
return nvcc_path
def install_cutt(root_folder):
# Modified from: https://github.com/ap-hynninen/cutt
url = "https://codeload.github.com/Jittor/cutt/zip/v1.2"
filename = "cutt-1.2.zip"
fullname = os.path.join(root_folder, filename)
dirname = os.path.join(root_folder, filename.replace(".zip",""))
true_md5 = "14d0fd1132c8cd657dc3cf29ce4db931"
if os.path.exists(fullname):
from jittor_utils.misc import calculate_md5
md5 = calculate_md5(fullname)
if md5 != true_md5:
os.remove(fullname)
shutil.rmtree(dirname)
if not os.path.isfile(os.path.join(cache_path, "libcutt"+so)):
LOG.i("Downloading cutt...")
download_url_to_local(url, filename, root_folder, true_md5)
import zipfile
zf = zipfile.ZipFile(fullname)
try:
zf.extractall(path=root_folder)
except RuntimeError as e:
print(e)
raise
zf.close()
LOG.i("installing cutt...")
arch_flag = ""
if len(flags.cuda_archs):
arch_flag = f" -arch=compute_{min(flags.cuda_archs)} "
arch_flag += ''.join(map(lambda x:f' -code=sm_{x} ', flags.cuda_archs))
cutt_include = f" -I\"{dirname}/include\" -I\"{dirname}/src\" "
files = glob.glob(dirname+"/src/*.c*", recursive=True)
files2 = []
for f in files:
if f.endswith("cutt_bench.cpp") or \
f.endswith("cutt_test.cpp"):
continue
files2.append(f)
cutt_flags = cc_flags+opt_flags+cutt_include
compile(cc_path, cutt_flags, files2, cache_path+"/libcutt"+so, cuda_flags=arch_flag)
return dirname
def download_url(self):
'''
Download mnist data set function, this function will be called when download is True.
'''
resources = [
("https://www.itl.nist.gov/iaui/vip/cs_links/EMNIST/gzip.zip", "58c8d27c78d21e728a6bc7b3cc06412e"),
]
for url, md5 in resources:
filename = "emnist.zip"
download_url_to_local(url, filename, self.data_root, md5)
import zipfile
zf = zipfile.ZipFile(os.path.join(self.data_root, filename))
try:
zf.extractall(path=self.data_root)
except RuntimeError as e:
print(e)
raise
zf.close()
def install_cub(root_folder):
url = "https://github.com/NVIDIA/cub/archive/1.11.0.tar.gz"
url = "https://codeload.github.com/NVIDIA/cub/tar.gz/1.11.0"
filename = "cub-1.11.0.tgz"
md5 = "97196a885598e40592100e1caaf3d5ea"
fullname = os.path.join(root_folder, filename)
dirname = os.path.join(root_folder, filename.replace(".tgz",""))
if not os.path.isfile(os.path.join(dirname, "examples", "device/example_device_radix_sort.cu")):
LOG.i("Downloading cub...")
download_url_to_local(url, filename, root_folder, md5)
import tarfile
with tarfile.open(fullname, "r") as tar:
tar.extractall(root_folder)
# assert 0 == os.system(f"cd {dirname}/examples && "
# f"{nvcc_path} --cudart=shared -ccbin=\"{cc_path}\" device/example_device_radix_sort.cu -O2 -I.. -std=c++14 -o test")
# if core.get_device_count():
# assert 0 == os.system(f"cd {dirname}/examples && ./test")
return dirname
def safeunpickle(path):
if path.startswith("jittorhub://"):
path = path.replace(
"jittorhub://",
"https://cg.cs.tsinghua.edu.cn/jittor/assets/build/checkpoints/")
if path.startswith("https:") or path.startswith("http:"):
base = path.split("/")[-1]
fname = os.path.join(compiler.ck_path, base)
from jittor_utils.misc import download_url_to_local
download_url_to_local(path, base, compiler.ck_path, None)
path = fname
if path.endswith(".pth"):
try:
dirty_fix_pytorch_runtime_error()
import torch
except:
raise RuntimeError(
"pytorch need to be installed when load pth format.")
model_dict = torch.load(path, map_location=torch.device('cpu'))
try:
for k, v in model_dict.items():
try:
if not isinstance(v, np.ndarray) and hasattr(v, "cpu"):
model_dict[k] = v.cpu().detach().numpy()
except:
pass
except:
pass
return model_dict
with open(path, "rb") as f:
s = f.read()
if not s.endswith(b"HCAJSLHD"):
return _load_pkl(s, path)
checksum = s[-28:-8]
s = s[:-28]
if hashlib.sha1(s).digest() != checksum:
raise ValueError(
"Pickle checksum does not match! path: " + path,
" This file maybe corrupted, please consider remove it"
" and re-download.")
return _load_pkl(s, path)
def install_nccl(root_folder):
url = "https://github.com/NVIDIA/nccl/archive/v2.8.4-1.tar.gz"
url = "https://codeload.github.com/NVIDIA/nccl/tar.gz/v2.8.4-1"
filename = "nccl.tgz"
fullname = os.path.join(root_folder, filename)
dirname = os.path.join(root_folder, "nccl-2.8.4-1")
true_md5 = "900666558c5bc43e0a5e84045b88a06f"
if os.path.exists(fullname):
md5 = run_cmd('md5sum ' + fullname).split()[0]
if md5 != true_md5:
os.remove(fullname)
if os.path.isdir(dirname):
shutil.rmtree(dirname)
if not os.path.isfile(os.path.join(dirname, "build", "lib", "libnccl.so")):
LOG.i("Downloading nccl...")
download_url_to_local(url, filename, root_folder, true_md5)
if core.get_device_count() == 0:
return
if not inside_mpi():
return
import tarfile
with tarfile.open(fullname, "r") as tar:
tar.extractall(root_folder)
LOG.i("installing nccl...")
arch_flag = ""
if len(flags.cuda_archs):
arch_flag = f" -arch=compute_{min(flags.cuda_archs)} "
arch_flag += ''.join(
map(lambda x: f' -code=sm_{x} ', flags.cuda_archs))
run_cmd(
f"CC=\"{cc_path}\" CXX=\"{cc_path}\" make -j8 src.build CUDA_HOME='{cuda_home}' NVCC_GENCODE='{arch_flag} --cudart=shared ' ",
cwd=dirname)
return dirname
def install_cutt(root_folder):
# Modified from: https://github.com/ap-hynninen/cutt
url = "https://codeload.github.com/Jittor/cutt/zip/v1.1"
filename = "cutt-1.1.zip"
fullname = os.path.join(root_folder, filename)
dirname = os.path.join(root_folder, filename.replace(".zip", ""))
true_md5 = "7bb71cf7c49dbe57772539bf043778f7"
if os.path.exists(fullname):
md5 = run_cmd('md5sum ' + fullname).split()[0]
if md5 != true_md5:
os.remove(fullname)
shutil.rmtree(dirname)
if not os.path.isfile(os.path.join(dirname, "bin", "cutt_test")):
LOG.i("Downloading cutt...")
download_url_to_local(url, filename, root_folder, true_md5)
import zipfile
zf = zipfile.ZipFile(fullname)
try:
zf.extractall(path=root_folder)
except RuntimeError as e:
print(e)
raise
zf.close()
LOG.i("installing cutt...")
arch_flag = ""
if len(flags.cuda_archs):
arch_flag = f" -arch=compute_{min(flags.cuda_archs)} "
arch_flag += ''.join(
map(lambda x: f' -code=sm_{x} ', flags.cuda_archs))
run_cmd(
f"make NVCC_GENCODE='{arch_flag} --cudart=shared -ccbin=\"{cc_path}\" ' nvcc_path='{nvcc_path}'",
cwd=dirname)
return dirname
def install_mkl(root_folder):
# origin url is
# url = "https://github.com/intel/mkl-dnn/releases/download/v1.0.2/mkldnn_lnx_1.0.2_cpu_gomp.tgz"
import platform
if platform.system() == "Linux":
if platform.machine() == 'x86_64':
filename = "dnnl_lnx_2.2.0_cpu_gomp.tgz"
md5 = "35bbbdf550a9d8ad54db798e372000f6"
elif platform.machine() == 'aarch64':
filename = "dnnl_lnx_2.2.0_cpu_gomp_aarch64.tgz"
md5 = "72cf9b0b8fd6c3c786d35a9daaee22b8"
else:
raise RuntimeError(
f"platform.machine()=={platform.machine()} not support yet,"
" Please contact us on https://github.com/jittor/jittor ")
else:
raise RuntimeError(
f"platform.machine()=={platform.machine()} not support yet,"
" Please contact us on https://github.com/jittor/jittor ")
url = "https://cg.cs.tsinghua.edu.cn/jittor/assets/" + filename
fullname = os.path.join(root_folder, filename)
dirname = os.path.join(root_folder, filename.replace(".tgz", ""))
if not os.path.isfile(os.path.join(dirname, "lib", "libmkldnn.so")):
LOG.i("Downloading mkl...")
download_url_to_local(url, filename, root_folder, md5)
import tarfile
with tarfile.open(fullname, "r") as tar:
tar.extractall(root_folder)
assert 0 == os.system(
f"cd {dirname}/examples && "
f"{cc_path} -std=c++14 cnn_inference_f32.cpp -Ofast -lmkldnn -I ../include -L ../lib -o test && LD_LIBRARY_PATH=../lib/ ./test"
)
def install_cuda():
cuda_driver_version = get_cuda_driver()
if not cuda_driver_version:
return None
LOG.i("cuda_driver_version: ", cuda_driver_version)
if "JTCUDA_VERSION" in os.environ:
cuda_driver_version = list(
map(int, os.environ["JTCUDA_VERSION"].split(".")))
LOG.i("JTCUDA_VERSION: ", cuda_driver_version)
if os.name == 'nt':
if cuda_driver_version >= [11, 4]:
cuda_tgz = "cuda11.4_cudnn8_win.zip"
md5 = "06eed370d0d44bb2cc57809343911187"
elif cuda_driver_version >= [11, 2]:
cuda_tgz = "cuda11.2_cudnn8_win.zip"
md5 = "b5543822c21bc460c1a414af47754556"
elif cuda_driver_version >= [
11,
]:
cuda_tgz = "cuda11.0_cudnn8_win.zip"
md5 = "7a248df76ee5e79623236b0560f8d1fd"
elif cuda_driver_version >= [
10,
]:
cuda_tgz = "cuda10.2_cudnn7_win.zip"
md5 = "7dd9963833a91371299a2ba58779dd71"
else:
raise RuntimeError(
f"Unsupport cuda driver version: {cuda_driver_version}, at least 10.2"
)
else:
if cuda_driver_version >= [11, 2]:
cuda_tgz = "cuda11.2_cudnn8_linux.tgz"
md5 = "b93a1a5d19098e93450ee080509e9836"
elif cuda_driver_version >= [
11,
]:
cuda_tgz = "cuda11.0_cudnn8_linux.tgz"
md5 = "5dbdb43e35b4db8249027997720bf1ca"
elif cuda_driver_version >= [10, 2]:
cuda_tgz = "cuda10.2_cudnn7_linux.tgz"
md5 = "40f0563e8eb176f53e55943f6d212ad7"
elif cuda_driver_version >= [
10,
]:
cuda_tgz = "cuda10.0_cudnn7_linux.tgz"
md5 = "f16d3ff63f081031d21faec3ec8b7dac"
else:
raise RuntimeError(
f"Unsupport cuda driver version: {cuda_driver_version}, at least 10.0"
)
jtcuda_path = os.path.join(pathlib.Path.home(), ".cache", "jittor",
"jtcuda")
nvcc_path = os.path.join(jtcuda_path, cuda_tgz[:-4], "bin", "nvcc")
if os.name == 'nt': nvcc_path += '.exe'
nvcc_lib_path = os.path.join(jtcuda_path, cuda_tgz[:-4], "lib64")
sys.path.append(nvcc_lib_path)
new_ld_path = os.environ.get("LD_LIBRARY_PATH", "") + ":" + nvcc_lib_path
os.environ["LD_LIBRARY_PATH"] = new_ld_path
if os.path.isfile(nvcc_path):
return nvcc_path
os.makedirs(jtcuda_path, exist_ok=True)
cuda_tgz_path = os.path.join(jtcuda_path, cuda_tgz)
download_url_to_local(
"https://cg.cs.tsinghua.edu.cn/jittor/assets/" + cuda_tgz, cuda_tgz,
jtcuda_path, md5)
if cuda_tgz.endswith(".zip"):
import zipfile
zf = zipfile.ZipFile(cuda_tgz_path)
zf.extractall(path=cuda_tgz_path[:-4])
else:
import tarfile
with tarfile.open(cuda_tgz_path, "r") as tar:
tar.extractall(cuda_tgz_path[:-4])
assert os.path.isfile(nvcc_path), nvcc_path
return nvcc_path