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)
