def _process(self): f = osp.join(self.processed_dir, 'pre_transform.pkl') if osp.exists(f) and jt.load(f) != __repr__(self.pre_transform): logging.warning( 'The `pre_transform` argument differs from the one used in ' 'the pre-processed version of this dataset. If you really ' 'want to make use of another pre-processing technique, make ' 'sure to delete `{}` first.'.format(self.processed_dir)) f = osp.join(self.processed_dir, 'pre_filter.pkl') if osp.exists(f) and jt.load(f) != __repr__(self.pre_filter): logging.warning( 'The `pre_filter` argument differs from the one used in the ' 'pre-processed version of this dataset. If you really want to ' 'make use of another pre-fitering technique, make sure to ' 'delete `{}` first.'.format(self.processed_dir)) if files_exist(self.processed_paths): # pragma: no cover return print('Processing...') makedirs(self.processed_dir) self.process() path = osp.join(self.processed_dir, 'pre_transform.pkl') jt.save(__repr__(self.pre_transform), path) path = osp.join(self.processed_dir, 'pre_filter.pkl') jt.save(__repr__(self.pre_filter), path) print('Done!')
def load_checkpoint(_model, resume_from, to_use_device, _optimizers=None, third_name=None): """ 加载预训练模型 Args: _model: 模型 resume_from: 预训练模型路径 to_use_device: 设备 _optimizers: 如果不为None,则表明采用模型的训练参数 third_name: 第三方预训练模型的名称 Returns: """ global_state = {} if not third_name: state = jittor.load(resume_from) _model.load_state_dict(state['state_dict']) if _optimizers is not None: _optimizers.load_state_dict(state['optimizer']) if 'global_state' in state: global_state = state['global_state'] elif third_name == 'paddle': import paddle.fluid as fluid paddle_model = fluid.io.load_program_state(resume_from) _model.load_3rd_state_dict(third_name, paddle_model) return _model, _optimizers, global_state
def load_deformable_vovnet(cfg, f): state_dict = jt.load(f) import re logger = logging.getLogger(__name__) logger.info("Remapping conv weights for deformable conv weights") layer_keys = sorted(state_dict.keys()) for idx, stage_with_dcn in enumerate(cfg.MODEL.VOVNET.STAGE_WITH_DCN, 2): if not stage_with_dcn: continue for old_key in layer_keys: if "layer" in old_key and f"OSA{idx}" in old_key and \ "conv" in old_key: for param in ["weight", "bias"]: if old_key.find(param) is -1: continue new_key = old_key.replace("conv.{}".format(param), "conv.conv.{}".format(param)) logger.info("old_key: {}, new_key: {}".format( old_key, new_key)) state_dict[new_key] = state_dict[old_key] del state_dict[old_key] if "model" not in state_dict: state_dict = dict(model=state_dict) return state_dict
def test_save(self): pp = [1,2,jt.array([1,2,3]), {"a":[1,2,3], "b":jt.array([1,2,3])}] jt.save(pp, "/tmp/xx.pkl") x = jt.load("/tmp/xx.pkl") assert x[:2] == [1,2] assert (x[2] == np.array([1,2,3])).all() assert x[3]['a'] == [1,2,3] assert (x[3]['b'] == np.array([1,2,3])).all()
def res2net101_v1b_26w_4s(pretrained=False, **kwargs): model = Res2Net(Bottle2neck, [3, 4, 23, 3], baseWidth=26, scale=4, **kwargs) if pretrained: model.load(jt.load((model_urls['res2net101_v1b_26w_4s']))) return model
def load_pretrained_model(model, params_path): pretrained_dict = jt.load(params_path) model_dict = {} param_name = model.parameters() name_list = [item.name() for item in param_name] for k, v in pretrained_dict.items(): if k in name_list: model_dict[k] = v model.load_parameters(model_dict)
def test_save(self): pp = [ 1, 2, jt.array([1, 2, 3]), { "a": [1, 2, 3], "b": jt.array([1, 2, 3]) } ] name = jt.flags.cache_path + "/xx.pkl" jt.save(pp, name) x = jt.load(name) assert x[:2] == [1, 2] assert (x[2] == np.array([1, 2, 3])).all() assert x[3]['a'] == [1, 2, 3] assert (x[3]['b'] == np.array([1, 2, 3])).all()
def main(): jittor.flags.use_cuda = 1 args = parse() task = args.task input_path, output_dir = args.input_path, args.output_dir sidelength, channels = args.sidelength, args.channels lr, epochs = args.lr, args.epochs if not os.path.exists(output_dir): os.mkdir(output_dir) network = model.Siren(channels) optim = jittor.optim.Adam(network.parameters(), lr=lr) network_path = os.path.join(output_dir, "network.pkl") coords, image_arr = dataio.get_fit_data(input_path, sidelength, channels) if task == 'fit': gt = image_arr.clone() loss_fn = compute.mse elif task == 'poisson': gt = dataio.get_poisson_data(input_path, sidelength, channels) loss_fn = compute.laplace_mse elif task == 'composition': gt = dataio.get_compos_data(input_path, args.ot_image_path, sidelength, channels) loss_fn = compute.two_grad_mse elif task == 'inpaint': data = dataio.get_inpaint_data(input_path, sidelength, channels, args.points) coords = data[SELECTED_COORD_KEY] loss_fn = compute.mse gt = data[SELECT_IMG_KEY] dataio.output_masked_image(output_dir, data[SELECT_IMG_KEY], data[IMAGE_KEY], data[MASK_KEY], sidelength, channels) else: raise ValueError() print("Training") train(network, optim, loss_fn, epochs, coords, gt, network_path) network.load_state_dict(jittor.load(network_path)) best_output = network(coords) if task == 'inpaint': best_output = network(data[COORD_KEY]) dataio.output_result(output_dir, best_output, coords, sidelength, channels, image_arr)
def __init__(self, root, name, split="public", num_train_per_class=20, num_val=500, num_test=1000, transform=None, pre_transform=None): self.name = name super(Planetoid, self).__init__(root, transform, pre_transform) self.data, self.slices = jt.load(self.processed_paths[0]) self.split = split assert self.split in ['public', 'full', 'random'] if split == 'full': data = self.get(0) init(data.train_mask, True) data.train_mask[jt.logical_or(data.val_mask, data.test_mask)] = False self.data, self.slices = self.collate([data]) elif split == 'random': data = self.get(0) init(data.train_mask, False) for c in range(self.num_classes): idx = (data.y == c).nonzero().view(-1) idx = idx[jt.randperm(idx.size(0))[:num_train_per_class]] data.train_mask[idx] = True remaining = jt.logical_not(data.train_mask).nonzero().view(-1) remaining = remaining[jt.randperm(remaining.size(0))] init(data.val_mask, False) data.val_mask[remaining[:num_val]] = True init(data.test_mask, False) data.test_mask[remaining[num_val:num_val + num_test]] = True self.data, self.slices = self.collate([data])
def _load_file(self, f): return jt.load(f)
default=128, help='size of the image') parser.add_argument('--n_row', type=int, default=5, help='number of rows of sample matrix') parser.add_argument('--n_col', type=int, default=5, help='number of columns of sample matrix') parser.add_argument('path', type=str, help='path to checkpoint file') args = parser.parse_args() generator = StyledGenerator(512) ckpt = jt.load(args.path) generator.load_state_dict(ckpt) generator.eval() mean_style = get_mean_style(generator) step = int(math.log(args.size, 2)) - 2 img = sample(generator, step, mean_style, args.n_row * args.n_col) jt.save_image(img, 'style_mixing/sample.png', nrow=args.n_col, normalize=True, range=(-1, 1)) for j in range(20):
np.random.seed(args.seed) directory_output = os.path.join(args.model_directory, args.experiment_id) os.makedirs(directory_output, exist_ok=True) image_output = os.path.join(directory_output, 'pic') os.makedirs(image_output, exist_ok=True) # setup model & optimizer model = models.Model('data/obj/sphere_642.obj', args=args) optimizer = nn.Adam(model.model_param(), args.learning_rate) start_iter = START_ITERATION if args.resume_path: state_dicts = jt.load(args.resume_path) model.load_state_dict(state_dicts['model']) optimizer.load_state_dict(state_dicts['optimizer']) start_iter = int(os.path.split(args.resume_path)[1][11:].split('.')[0]) + 1 print('Resuming from %s iteration' % start_iter) dataset_train = datasets.ShapeNet(args.dataset_directory, args.class_ids.split(','), 'train') def train(): end = time.time() batch_time = AverageMeter() data_time = AverageMeter() losses = AverageMeter()
g_running.eval() d_optimizer = jt.optim.Adam(netD.parameters(), lr=lr, betas=(0.0, 0.99)) g_optimizer = jt.optim.Adam(netG.generator.parameters(), lr=lr, betas=(0.0, 0.99)) g_optimizer.add_param_group({ 'params': netG.style.parameters(), 'lr': lr * 0.01, 'mult': 0.01, }) accumulate(g_running, netG, 0) if args.ckpt is not None: ckpt = jt.load(args.ckpt) netG.load_state_dict(ckpt['generator']) netD.load_state_dict(ckpt['discriminator']) g_running.load_state_dict(ckpt['g_running']) print('resuming from checkpoint .......') ## Actual Training step = init_step resolution = int(4 * 2**step) image_loader = SymbolDataset(args.path, transform, resolution).set_attrs( batch_size=batch_size.get(resolution, batch_default), shuffle=True) train_loader = iter(image_loader) requires_grad(netG, False)