# All Rights Reserved. # This file is subject to the terms and conditions defined in # file 'LICENSE.txt', which is part of this source code package. # *************************************************************** __version__ = '1.1.5.6' from . import lock with lock.lock_scope(): from . import compiler from .compiler import LOG, has_cuda from .compiler import compile_custom_ops, compile_custom_op import jittor_core as core from jittor_core import * from jittor_core.ops import * from . import compile_extern from .compile_extern import mkl_ops, mpi, mpi_ops, in_mpi, rank if core.get_device_count() == 0: has_cuda = compile_extern.has_cuda = compiler.has_cuda = False if has_cuda: from .compile_extern import cudnn, curand, cublas import contextlib import numpy as np from collections import OrderedDict from collections.abc import Sequence, Mapping import types import pickle import sys import traceback class _call_no_record_scope: def __enter__(self): pass
import numpy as np import cv2 from PIL import Image import time from util.Visualizer import Visualizer import jittor as jt from collections import OrderedDict from options.train_options import TrainOptions import util.util as util from models.IS_Model import IS_Model from jittor.dataset import Dataset import jittor_core import os jt.flags.use_cuda = 1 os.environ["CUDA_VISIBLE_DEVICES"] = "1,2" print(jittor_core.get_device_count()) def make_dataset(dir): images = [] for root, _, fnames in sorted(os.walk(dir)): for fname in fnames: path = os.path.join(root, fname) images.append(path) return images class AlignedDataset(Dataset): def initialize(self, root): self.part = { 'bg': (0, 0, 512),