def __init__(self, policy, policy_target, qfns, qfns_target, sampler, *,
FLAGS):
super().__init__()
self.FLAGS = FLAGS
self.policy = policy
self.policy_target = policy_target
self.qfns = qfns
self.qfns_target = qfns_target
self.sampler = sampler
self.qfns_opt = nn.Adam(
sum([qfn.parameters() for qfn in self.qfns], []), FLAGS.lr)
self.policy_opt = nn.Adam(self.policy.parameters(), FLAGS.lr)
self.n_batches = 0
def main():
parser = argparse.ArgumentParser()
parser.add_argument('-io', '--filename_obj', type=str, default=os.path.join(data_dir, 'teapot.obj'))
parser.add_argument('-ir', '--filename_ref', type=str, default=os.path.join(data_dir, 'example4_ref.png'))
parser.add_argument('-or', '--filename_output', type=str, default=os.path.join(data_dir, 'example4_result.gif'))
parser.add_argument('-mr', '--make_reference_image', type=int, default=0)
parser.add_argument('-g', '--gpu', type=int, default=0)
args = parser.parse_args()
if args.make_reference_image:
make_reference_image(args.filename_ref, args.filename_obj)
model = Model(args.filename_obj, args.filename_ref)
optimizer = nn.Adam(model.parameters(), lr=0.1)
loop = tqdm.tqdm(range(1000))
for i in loop:
loss = model()
optimizer.step(loss)
images, _, _ = model.renderer(model.vertices, model.faces, jt.tanh(model.textures))
image = images.numpy()[0].transpose(1,2,0)
imsave('/tmp/_tmp_%04d.png' % i, image)
loop.set_description('Optimizing (loss %.4f)' % loss.data)
if loss.data < 70:
break
make_gif(args.filename_output)
def main():
parser = argparse.ArgumentParser()
parser.add_argument('-i', '--filename-input', type=str,
default=os.path.join(data_dir, 'source.npy'))
parser.add_argument('-c', '--camera-input', type=str,
default=os.path.join(data_dir, 'camera.npy'))
parser.add_argument('-t', '--template-mesh', type=str,
default=os.path.join(data_dir, 'obj/sphere/sphere_1352.obj'))
parser.add_argument('-o', '--output-dir', type=str,
default=os.path.join(data_dir, 'results/output_deform'))
parser.add_argument('-b', '--batch-size', type=int,
default=120)
args = parser.parse_args()
os.makedirs(args.output_dir, exist_ok=True)
model = Model(args.template_mesh)
renderer = jr.Renderer(image_size=64, sigma_val=1e-4, aggr_func_rgb='hard', camera_mode='look_at', viewing_angle=15, dr_type='softras')
# read training images and camera poses
images = np.load(args.filename_input).astype('float32') / 255.
cameras = np.load(args.camera_input).astype('float32')
optimizer = nn.Adam(model.parameters(), 0.01, betas=(0.5, 0.99))
camera_distances = jt.array(cameras[:, 0])
elevations = jt.array(cameras[:, 1])
viewpoints = jt.array(cameras[:, 2])
renderer.transform.set_eyes_from_angles(camera_distances, elevations, viewpoints)
import time
sta = time.time()
loop = tqdm.tqdm(list(range(0, 1000)))
writer = imageio.get_writer(os.path.join(args.output_dir, 'deform.gif'), mode='I')
for i in loop:
images_gt = jt.array(images)
mesh, laplacian_loss, flatten_loss = model(args.batch_size)
images_pred = renderer.render_mesh(mesh, mode='silhouettes')
# optimize mesh with silhouette reprojection error and
# geometry constraints
loss = neg_iou_loss(images_pred, images_gt[:, 3]) + \
0.03 * laplacian_loss + \
0.0003 * flatten_loss
loop.set_description('Loss: %.4f'%(loss.item()))
optimizer.step(loss)
if i % 100 == 0:
image = images_pred.numpy()[0]#.transpose((1, 2, 0))
imageio.imsave(os.path.join(args.output_dir, 'deform_%05d.png'%i), (255*image).astype(np.uint8))
writer.append_data((255*image).astype(np.uint8))
# save optimized mesh
model(1)[0].save_obj(os.path.join(args.output_dir, 'plane.obj'), save_texture=False)
print(f"Cost {time.time() - sta} secs.")
def main():
parser = argparse.ArgumentParser()
parser.add_argument('-io',
'--filename_obj',
type=str,
default=os.path.join(data_dir, 'teapot.obj'))
parser.add_argument('-ir',
'--filename_ref',
type=str,
default=os.path.join(data_dir, 'example2_ref.png'))
parser.add_argument('-oo',
'--filename_output_optimization',
type=str,
default=os.path.join(data_dir,
'example2_optimization.gif'))
parser.add_argument('-or',
'--filename_output_result',
type=str,
default=os.path.join(data_dir, 'example2_result.gif'))
parser.add_argument('-g', '--gpu', type=int, default=0)
args = parser.parse_args()
model = Model(args.filename_obj, args.filename_ref)
optimizer = nn.Adam(model.parameters(), lr=1e-3)
loop = tqdm.tqdm(range(300))
for i in loop:
loop.set_description('Optimizing')
loss = model()
optimizer.step(loss)
images = model.renderer(model.vertices,
model.faces,
mode='silhouettes')
image = images.numpy()[0, 0]
imsave('/tmp/_tmp_%04d.png' % i, image)
make_gif(args.filename_output_optimization)
# draw object
loop = tqdm.tqdm(range(0, 360, 4))
for num, azimuth in enumerate(loop):
loop.set_description('Drawing')
model.renderer.eye = nr.get_points_from_angles(2.732, 0, azimuth)
images, _, _ = model.renderer(model.vertices, model.faces,
model.textures)
image = images.numpy()[0].transpose((1, 2, 0))
imsave('/tmp/_tmp_%04d.png' % num, image)
make_gif(args.filename_output_result)
def main():
parser = argparse.ArgumentParser()
parser.add_argument('-io',
'--filename_obj',
type=str,
default=os.path.join(data_dir,
'obj/spot/spot_triangulated.obj'))
parser.add_argument('-ir',
'--filename_ref',
type=str,
default=os.path.join(data_dir,
'ref/ref_roughness.png'))
parser.add_argument('-or',
'--filename_output',
type=str,
default=os.path.join(
data_dir,
'results/output_optim_roughness_textures'))
parser.add_argument('-g', '--gpu', type=int, default=0)
args = parser.parse_args()
os.makedirs(args.filename_output, exist_ok=True)
model = Model(args.filename_obj, args.filename_ref)
optimizer = nn.Adam([model.roughness_textures], lr=0.1, betas=(0.5, 0.999))
loop = tqdm.tqdm(range(15))
for num in loop:
loop.set_description('Optimizing')
loss = model()
optimizer.step(loss)
model.renderer.transform.set_eyes_from_angles(2.732, 30, 140)
images = model.renderer(model.vertices,
model.faces,
model.textures,
metallic_textures=model.metallic_textures,
roughness_textures=model.roughness_textures)
image = images.numpy()[0].transpose((1, 2, 0))
imsave('/tmp/_tmp_%04d.png' % num, image)
make_gif(os.path.join(args.filename_output, 'result.gif'))
discriminator = Discriminator()
# 导入MNIST数据集
from jittor.dataset.mnist import MNIST
import jittor.transform as transform
transform = transform.Compose([
transform.Resize(opt.img_size),
transform.Gray(),
transform.ImageNormalize(mean=[0.5], std=[0.5]),
])
dataloader = MNIST(train=True,
transform=transform).set_attrs(batch_size=opt.batch_size,
shuffle=True)
optimizer_G = nn.Adam(generator.parameters(),
lr=opt.lr,
betas=(opt.b1, opt.b2))
optimizer_D = nn.Adam(discriminator.parameters(),
lr=opt.lr,
betas=(opt.b1, opt.b2))
from PIL import Image
def save_image(img, path, nrow=10, padding=5):
N, C, W, H = img.shape
if (N % nrow != 0):
print("N%nrow!=0")
return
ncol = int(N / nrow)
img_all = []
output = net(pts, normals)
pred = np.argmax(output.data, axis=1)
acc = np.sum(pred == labels.data)
total_acc += acc
total_num += labels.shape[0]
acc = total_acc / total_num
return acc
if __name__ == '__main__':
freeze_random_seed()
net = PointNet(n_classes=40)
optimizer = nn.Adam(net.parameters(), lr=1e-3)
lr_scheduler = LRScheduler(optimizer)
batch_size = 32
train_dataloader = ModelNet40(n_points=4096, batch_size=batch_size, train=True, shuffle=True)
val_dataloader = ModelNet40(n_points=4096, batch_size=batch_size, train=False, shuffle=False)
step = 0
best_acc = 0
for epoch in range(1000):
lr_scheduler.step(len(train_dataloader) * batch_size)
train(net, optimizer, epoch, train_dataloader)
acc = evaluate(net, epoch, val_dataloader)
for conv in self.convs:
x = nn.dropout(x, self.dropout)
x = conv(x, x_0, edge_index, edge_weight)
x = nn.relu(x)
x = nn.dropout(x, self.dropout)
x = self.lins[1](x)
return nn.log_softmax(x, dim=-1)
model = Net(hidden_channels=64, num_layers=64, alpha=0.1, theta=0.5,
shared_weights=True, dropout=0.6)
optimizer = nn.Adam([
dict(params=model.convs.parameters(), weight_decay=0.01),
dict(params=model.lins.parameters(), weight_decay=5e-4)
], lr=0.01)
def train():
model.train()
out = model()[data.train_mask]
label = data.y[data.train_mask]
loss = nn.nll_loss(
out, label)
optimizer.step(loss)
return float(loss)
def test():
model.eval()
def __init__(self):
super(Net, self).__init__()
self.conv1 = SGConv(dataset.num_features,
dataset.num_classes,
K=2,
cached=True)
def execute(self):
x, edge_index = data.x, data.edge_index
x = self.conv1(x, edge_index)
return nn.log_softmax(x, dim=1)
model, data = Net(), data
# Only perform weight-decay on first convolution.
optimizer = nn.Adam(model.parameters(), lr=0.2, weight_decay=0.005)
def train():
model.train()
pred = model()[data.train_mask]
label = data.y[data.train_mask]
loss = nn.nll_loss(pred, label)
# print(loss)
optimizer.step(loss)
def test():
model.eval()
logits, accs = model(), []
for _, mask in data('train_mask', 'val_mask', 'test_mask'):
coupled_generators = CoupledGenerators()
coupled_discriminators = CoupledDiscriminators()
print(coupled_generators)
print(coupled_discriminators)
transform = transform.Compose([
transform.Resize(opt.img_size),
transform.ImageNormalize(mean=[0.5], std=[0.5]),
])
dataloader1 = MNIST(train=True, transform=transform).set_attrs(batch_size=opt.batch_size, shuffle=True)
dataloader2 = mnistm.MNISTM(mnist_root = "../../data/mnistm", train=True, transform = transform).set_attrs(batch_size=opt.batch_size, shuffle=True)
# Optimizers
optimizer_G = nn.Adam(coupled_generators.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))
optimizer_D = nn.Adam(coupled_discriminators.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))
# ----------
# Training
# ----------
for epoch in range(opt.n_epochs):
for i, ((imgs1, _), (imgs2, _)) in enumerate(zip(dataloader1, dataloader2)):
jt.sync_all(True)
batch_size = imgs1.shape[0]
# Adversarial ground truths
valid = jt.ones([batch_size, 1]).float32().stop_grad()
fake = jt.zeros([batch_size, 1]).float32().stop_grad()
# ------------------
# Configure data loader
from jittor.dataset.mnist import MNIST
import jittor.transform as transform
transform = transform.Compose([
transform.Resize(opt.img_size),
transform.Gray(),
transform.ImageNormalize(mean=[0.5], std=[0.5]),
])
train_loader = MNIST(train=True,
transform=transform).set_attrs(batch_size=opt.batch_size,
shuffle=True)
# Optimizers
optimizer_G = nn.Adam(encoder.parameters() + decoder.parameters(),
lr=opt.lr,
betas=(opt.b1, opt.b2))
optimizer_D = nn.Adam(discriminator.parameters(),
lr=opt.lr,
betas=(opt.b1, opt.b2))
def save_image(img, path, nrow=10, padding=5):
N, C, W, H = img.shape
if (N % nrow != 0):
print("N%nrow!=0")
return
ncol = int(N / nrow)
img_all = []
for i in range(ncol):
img_ = []
def train_net(cfg):
# Enable the inbuilt cudnn auto-tuner to find the best algorithm to use
# train_dataset_loader = utils.data_loaders.DATASET_LOADER_MAPPING[cfg.DATASET.TRAIN_DATASET](cfg)
# test_dataset_loader = utils.data_loaders.DATASET_LOADER_MAPPING[cfg.DATASET.TEST_DATASET](cfg)
train_dataset_loader = dataloader_jt.DATASET_LOADER_MAPPING[
cfg.DATASET.TRAIN_DATASET](cfg)
test_dataset_loader = dataloader_jt.DATASET_LOADER_MAPPING[
cfg.DATASET.TEST_DATASET](cfg)
train_data_loader = train_dataset_loader.get_dataset(
dataloader_jt.DatasetSubset.TRAIN,
batch_size=cfg.TRAIN.BATCH_SIZE,
num_workers=cfg.CONST.NUM_WORKERS,
shuffle=True)
val_data_loader = test_dataset_loader.get_dataset(
dataloader_jt.DatasetSubset.VAL,
batch_size=cfg.TRAIN.BATCH_SIZE,
num_workers=cfg.CONST.NUM_WORKERS,
shuffle=False)
# Set up folders for logs and checkpoints
output_dir = os.path.join(cfg.DIR.OUT_PATH, '%s',
datetime.now().isoformat())
cfg.DIR.CHECKPOINTS = output_dir % 'checkpoints'
cfg.DIR.LOGS = output_dir % 'logs'
if not os.path.exists(cfg.DIR.CHECKPOINTS):
os.makedirs(cfg.DIR.CHECKPOINTS)
# Create tensorboard writers
train_writer = SummaryWriter(os.path.join(cfg.DIR.LOGS, 'train'))
val_writer = SummaryWriter(os.path.join(cfg.DIR.LOGS, 'test'))
model = Model(dataset=cfg.DATASET.TRAIN_DATASET)
init_epoch = 0
best_metrics = float('inf')
optimizer = nn.Adam(model.parameters(),
lr=cfg.TRAIN.LEARNING_RATE,
weight_decay=cfg.TRAIN.WEIGHT_DECAY,
betas=cfg.TRAIN.BETAS)
lr_scheduler = jittor.lr_scheduler.MultiStepLR(
optimizer,
milestones=cfg.TRAIN.LR_MILESTONES,
gamma=cfg.TRAIN.GAMMA,
last_epoch=init_epoch)
# Training/Testing the network
for epoch_idx in range(init_epoch + 1, cfg.TRAIN.N_EPOCHS + 1):
epoch_start_time = time()
model.train()
loss_metric = AverageMeter()
n_batches = len(train_data_loader)
print('epoch: ', epoch_idx, 'optimizer: ', lr_scheduler.get_lr())
with tqdm(train_data_loader) as t:
for batch_idx, (taxonomy_ids, model_ids, data) in enumerate(t):
partial = jittor.array(data['partial_cloud'])
gt = jittor.array(data['gtcloud'])
pcds, deltas = model(partial)
cd1 = chamfer(pcds[0], gt)
cd2 = chamfer(pcds[1], gt)
cd3 = chamfer(pcds[2], gt)
loss_cd = cd1 + cd2 + cd3
delta_losses = []
for delta in deltas:
delta_losses.append(jittor.sum(delta**2))
loss_pmd = jittor.sum(jittor.stack(delta_losses)) / 3
loss = loss_cd * cfg.TRAIN.LAMBDA_CD + loss_pmd * cfg.TRAIN.LAMBDA_PMD
optimizer.step(loss)
loss_item = loss.item()
loss_metric.update(loss_item)
jittor.sync_all()
t.set_description(
'[Epoch %d/%d][Batch %d/%d]' %
(epoch_idx, cfg.TRAIN.N_EPOCHS, batch_idx + 1, n_batches))
t.set_postfix(loss='%s' % ['%.4f' % l for l in [loss_item]])
lr_scheduler.step()
epoch_end_time = time()
train_writer.add_scalar('Loss/Epoch/loss', loss_metric.avg(),
epoch_idx)
logging.info(
'[Epoch %d/%d] EpochTime = %.3f (s) Losses = %s' %
(epoch_idx, cfg.TRAIN.N_EPOCHS, epoch_end_time - epoch_start_time,
['%.4f' % l for l in [loss_metric.avg()]]))
# Validate the current model
cd_eval = test_net(cfg, epoch_idx, val_data_loader, val_writer, model)
# Save checkpoints
if epoch_idx % cfg.TRAIN.SAVE_FREQ == 0 or cd_eval < best_metrics:
file_name = 'ckpt-best.pkl' if cd_eval < best_metrics else 'ckpt-epoch-%03d.pkl' % epoch_idx
output_path = os.path.join(cfg.DIR.CHECKPOINTS, file_name)
model.save(output_path)
logging.info('Saved checkpoint to %s ...' % output_path)
if cd_eval < best_metrics:
best_metrics = cd_eval
train_writer.close()
val_writer.close()
def __setup_dis_optim(self, learning_rate, beta_1, beta_2, eps):
# self.dis_optim = torch.optim.Adam(self.dis.parameters(), lr=learning_rate, betas=(beta_1, beta_2), eps=eps)
self.dis_optim = nn.Adam(self.dis.parameters(),
lr=learning_rate,
betas=(beta_1, beta_2),
eps=eps)
dataset.num_classes,
cached=True,
normalize=not args.use_gdc)
def execute(self):
x, edge_index, edge_weight = data.x, data.edge_index, data.edge_attr
x = nn.relu(self.conv1(x, edge_index, edge_weight))
x = nn.dropout(x)
x = self.conv2(x, edge_index, edge_weight)
return nn.log_softmax(x, dim=1)
model, data = Net(), data
optimizer = nn.Adam([
dict(params=model.conv1.parameters(), weight_decay=5e-4),
dict(params=model.conv2.parameters(), weight_decay=0)
],
lr=0.01) # Only perform weight-decay on first convolution.
def train():
model.train()
pred = model()[data.train_mask]
label = data.y[data.train_mask]
loss = nn.nll_loss(pred, label)
optimizer.step(loss)
def test():
model.eval()
logits, accs = model(), []
parser.add_argument('-df', '--demo-freq', type=int, default=DEMO_FREQ)
parser.add_argument('-sf', '--save-freq', type=int, default=SAVE_FREQ)
parser.add_argument('-s', '--seed', type=int, default=RANDOM_SEED)
args = parser.parse_args()
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()
