def test_resnet_infer_with_feature(self): cat_url = "https://ss1.bdstatic.com/70cFuXSh_Q1YnxGkpoWK1HF6hhy/it/u=3782485413,1118109468&fm=26&gp=0.jpg" import jittor_utils cat_path = f"{jt.flags.cache_path}/cat.jpg" print("download") jittor_utils.download(cat_url, cat_path) with open(cat_path, 'rb') as f: img = Image.open(f).convert('RGB') img = jt.array(np.array(img)) print(img.shape, img.dtype) img = ((img.float() - 128) / 255).transpose(2, 0, 1) with jt.flag_scope(trace_py_var=2, trace_var_data=1): img = img[None, ...] resnet18 = resnet.Resnet18(pretrained=True) x = jt.float32(img) y = resnet18(x) y.sync() data = jt.dump_trace_data() jt.clear_trace_data() with open(f"{jt.flags.cache_path}/resnet_with_feature.pkl", "wb") as f: pickle.dump(data, f) for k, v in data["execute_op_info"].items(): for i in v['fused_ops']: if i not in data["node_data"]: assert 0, (i, "not found")
# os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" libname = {"clang": "omp", "icc": "iomp5", "g++": "gomp"}[cc_type] libname = ctypes.util.find_library(libname) assert libname is not None, "openmp library not found" ctypes.CDLL(libname, os.RTLD_NOW | os.RTLD_GLOBAL) version_file = os.path.join(jittor_path, "version") if os.path.isfile(version_file): with open(version_file, 'r') as f: version = f.read().strip() key = f"{version}-{cc_type}-{'cuda' if has_cuda else 'cpu'}.o" # TODO: open the website extra_obj = os.path.join(cache_path, key) url = os.path.join("https://cg.cs.tsinghua.edu.cn/jittor/assets/build/" + key) jit_utils.download(url, extra_obj) files.append(extra_obj) compile(cc_path, cc_flags + opt_flags, files, 'jittor_core' + extension_suffix) # TODO: move to compile_extern.py compile_extern() with jit_utils.import_scope(import_flags): import jittor_core as core flags = core.flags() if has_cuda: nvcc_flags += f" -arch={','.join(map(lambda x:'sm_'+str(x),flags.cuda_archs))} " flags.cc_path = cc_path