def main(cfg: DictConfig):
# Set the relevant seeds for reproducibility.
np.random.seed(cfg.seed)
torch.manual_seed(cfg.seed)
# Device on which to run.
if torch.cuda.is_available():
device = "cuda"
else:
warnings.warn(
"Please note that although executing on CPU is supported," +
"the training is unlikely to finish in resonable time.")
device = "cpu"
# Initialize the Radiance Field model.
model = RadianceFieldRenderer(
image_size=cfg.data.image_size,
n_pts_per_ray=cfg.raysampler.n_pts_per_ray,
n_pts_per_ray_fine=cfg.raysampler.n_pts_per_ray,
n_rays_per_image=cfg.raysampler.n_rays_per_image,
min_depth=cfg.raysampler.min_depth,
max_depth=cfg.raysampler.max_depth,
stratified=cfg.raysampler.stratified,
stratified_test=cfg.raysampler.stratified_test,
chunk_size_test=cfg.raysampler.chunk_size_test,
n_harmonic_functions_xyz=cfg.implicit_function.
n_harmonic_functions_xyz,
n_harmonic_functions_dir=cfg.implicit_function.
n_harmonic_functions_dir,
n_hidden_neurons_xyz=cfg.implicit_function.n_hidden_neurons_xyz,
n_hidden_neurons_dir=cfg.implicit_function.n_hidden_neurons_dir,
n_layers_xyz=cfg.implicit_function.n_layers_xyz,
density_noise_std=cfg.implicit_function.density_noise_std,
)
# Move the model to the relevant device.
model.to(device)
# Init stats to None before loading.
stats = None
optimizer_state_dict = None
start_epoch = 0
checkpoint_path = os.path.join(hydra.utils.get_original_cwd(),
cfg.checkpoint_path)
if len(cfg.checkpoint_path) > 0:
# Make the root of the experiment directory.
checkpoint_dir = os.path.split(checkpoint_path)[0]
os.makedirs(checkpoint_dir, exist_ok=True)
# Resume training if requested.
if cfg.resume and os.path.isfile(checkpoint_path):
print(f"Resuming from checkpoint {checkpoint_path}.")
loaded_data = torch.load(checkpoint_path)
model.load_state_dict(loaded_data["model"])
stats = pickle.loads(loaded_data["stats"])
print(f" => resuming from epoch {stats.epoch}.")
optimizer_state_dict = loaded_data["optimizer"]
start_epoch = stats.epoch
# Initialize the optimizer.
optimizer = torch.optim.Adam(
model.parameters(),
lr=cfg.optimizer.lr,
)
# Load the optimizer state dict in case we are resuming.
if optimizer_state_dict is not None:
optimizer.load_state_dict(optimizer_state_dict)
optimizer.last_epoch = start_epoch
# Init the stats object.
if stats is None:
stats = Stats([
"loss", "mse_coarse", "mse_fine", "psnr_coarse", "psnr_fine",
"sec/it"
], )
# Learning rate scheduler setup.
# Following the original code, we use exponential decay of the
# learning rate: current_lr = base_lr * gamma ** (epoch / step_size)
def lr_lambda(epoch):
return cfg.optimizer.lr_scheduler_gamma**(
epoch / cfg.optimizer.lr_scheduler_step_size)
# The learning rate scheduling is implemented with LambdaLR PyTorch scheduler.
lr_scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer,
lr_lambda,
last_epoch=start_epoch -
1,
verbose=False)
# Initialize the cache for storing variables needed for visulization.
visuals_cache = collections.deque(maxlen=cfg.visualization.history_size)
# Init the visualization visdom env.
if cfg.visualization.visdom:
viz = Visdom(
server=cfg.visualization.visdom_server,
port=cfg.visualization.visdom_port,
use_incoming_socket=False,
)
else:
viz = None
# Load the training/validation data.
train_dataset, val_dataset, _ = get_nerf_datasets(
dataset_name=cfg.data.dataset_name,
image_size=cfg.data.image_size,
)
if cfg.data.precache_rays:
# Precache the projection rays.
model.eval()
with torch.no_grad():
for dataset in (train_dataset, val_dataset):
cache_cameras = [e["camera"].to(device) for e in dataset]
cache_camera_hashes = [e["camera_idx"] for e in dataset]
model.precache_rays(cache_cameras, cache_camera_hashes)
train_dataloader = torch.utils.data.DataLoader(
train_dataset,
batch_size=1,
shuffle=True,
num_workers=0,
collate_fn=trivial_collate,
)
# The validation dataloader is just an endless stream of random samples.
val_dataloader = torch.utils.data.DataLoader(
val_dataset,
batch_size=1,
num_workers=0,
collate_fn=trivial_collate,
sampler=torch.utils.data.RandomSampler(
val_dataset,
replacement=True,
num_samples=cfg.optimizer.max_epochs,
),
)
# Set the model to the training mode.
model.train()
# Run the main training loop.
for epoch in range(start_epoch, cfg.optimizer.max_epochs):
stats.new_epoch() # Init a new epoch.
for iteration, batch in enumerate(train_dataloader):
image, camera, camera_idx = batch[0].values()
image = image.to(device)
camera = camera.to(device)
optimizer.zero_grad()
# Run the forward pass of the model.
nerf_out, metrics = model(
camera_idx if cfg.data.precache_rays else None,
camera,
image,
)
# The loss is a sum of coarse and fine MSEs
loss = metrics["mse_coarse"] + metrics["mse_fine"]
# Take the training step.
loss.backward()
optimizer.step()
# Update stats with the current metrics.
stats.update(
{
"loss": float(loss),
**metrics
},
stat_set="train",
)
if iteration % cfg.stats_print_interval == 0:
stats.print(stat_set="train")
# Update the visualisatioon cache.
visuals_cache.append({
"camera":
camera.cpu(),
"camera_idx":
camera_idx,
"image":
image.cpu().detach(),
"rgb_fine":
nerf_out["rgb_fine"].cpu().detach(),
"rgb_coarse":
nerf_out["rgb_coarse"].cpu().detach(),
"rgb_gt":
nerf_out["rgb_gt"].cpu().detach(),
"coarse_ray_bundle":
nerf_out["coarse_ray_bundle"],
})
# Adjust the learning rate.
lr_scheduler.step()
print(cfg.validation_epoch_interval)
# Validation
if epoch % cfg.validation_epoch_interval == 0 and epoch > 0:
# Sample a validation camera/image.
val_batch = next(val_dataloader.__iter__())
val_image, val_camera, camera_idx = val_batch[0].values()
val_image = val_image.to(device)
val_camera = val_camera.to(device)
# Activate eval mode of the model (allows to do a full rendering pass).
model.eval()
with torch.no_grad():
val_nerf_out, val_metrics = model(
camera_idx if cfg.data.precache_rays else None,
val_camera,
val_image,
)
# Update stats with the validation metrics.
stats.update(val_metrics, stat_set="val")
stats.print(stat_set="val")
if viz is not None:
# Plot that loss curves into visdom.
stats.plot_stats(
viz=viz,
visdom_env=cfg.visualization.visdom_env,
plot_file=None,
)
# Visualize the intermediate results.
visualize_nerf_outputs(val_nerf_out, visuals_cache, viz,
cfg.visualization.visdom_env)
# Set the model back to train mode.
model.train()
# Checkpoint.
if (epoch % cfg.checkpoint_epoch_interval == 0
and len(cfg.checkpoint_path) > 0 and epoch > 0):
print(f"Storing checkpoint {checkpoint_path}.")
data_to_store = {
"model": model.state_dict(),
"optimizer": optimizer.state_dict(),
"stats": pickle.dumps(stats),
}
torch.save(data_to_store, checkpoint_path)
def main(cfg: DictConfig):
np.random.seed(cfg.seed)
torch.manual_seed(cfg.seed)
obj_path = cfg.data.obj_path
texture_path = cfg.data.texture_path
views_folder = cfg.data.views_folder
params_file = os.path.join(views_folder,"params.json")
dataset = CowMultiViews(obj_path,views_folder,texture_path,params_file=params_file)
train_dataset, validation_dataset, test_dataset = CowMultiViews.random_split_dataset(dataset,
train_fraction=0.7,
validation_fraction=0.2)
del dataset
train_dataset.unit_normalize()
validation_dataset.unit_normalize()
mesh_verts = train_dataset.get_verts()
mesh_edges = train_dataset.get_edges()
mesh_vert_normals = train_dataset.get_vert_normals()
mesh_texture = train_dataset.get_texture()
pytorch_mesh = train_dataset.pytorch_mesh.cuda()
random_face_attrs = train_dataset.get_faces_as_vertex_matrices(features_list=['random'],num_random_dims=cfg.training.feature_dim)
coord_face_attrs = train_dataset.get_faces_as_vertex_matrices(features_list=['coord'],num_random_dims=cfg.training.feature_dim)
normal_face_attrs = train_dataset.get_faces_as_vertex_matrices(features_list=['normal'],num_random_dims=cfg.training.feature_dim)
torch_verts = torch.from_numpy(np.array(mesh_verts)).float().cuda()
torch_edges = torch.from_numpy(np.array(mesh_edges)).long().cuda()
torch_normals = torch.from_numpy(np.array(mesh_vert_normals)).float().cuda()
torch_texture = torch.from_numpy(np.array(mesh_texture)).float().cuda()
torch_texture = torch.unsqueeze(torch_texture,0)
torch_random_face_attrs = torch.tensor(np.array(random_face_attrs),requires_grad=True).float().cuda()
torch_random_face_attrs = torch.nn.Parameter(torch_random_face_attrs)
torch_coord_face_attrs = torch.tensor(np.array(coord_face_attrs)).float().cuda()
torch_normal_face_attrs = torch.tensor(np.array(normal_face_attrs)).float().cuda()
train_dataloader = DataLoader(train_dataset,batch_size=cfg.training.batch_size,shuffle=True,num_workers=4)
validation_dataloader = DataLoader(validation_dataset,batch_size=cfg.training.batch_size,shuffle=True,num_workers=4)
image_translator = ImageTranslator(input_dim=cfg.training.feature_dim+9,output_dim=3,
image_size=tuple(cfg.data.image_size)).cuda()
mse_loss = torch.nn.MSELoss()
# Initialize the optimizer.
optimizer = torch.optim.Adam(
list(image_translator.parameters())+[torch_random_face_attrs],
lr=cfg.optimizer.lr,
)
stats = None
start_epoch = 0
checkpoint_path = os.path.join(hydra.utils.get_original_cwd(), cfg.checkpoint_path)
# Init the stats object.
if stats is None:
stats = Stats(
["mse_loss", "sec/it"],
)
# Learning rate scheduler setup.
# Following the original code, we use exponential decay of the
# learning rate: current_lr = base_lr * gamma ** (epoch / step_size)
def lr_lambda(epoch):
return cfg.optimizer.lr_scheduler_gamma ** (
epoch / cfg.optimizer.lr_scheduler_step_size
)
# The learning rate scheduling is implemented with LambdaLR PyTorch scheduler.
lr_scheduler = torch.optim.lr_scheduler.LambdaLR(
optimizer, lr_lambda, last_epoch=start_epoch - 1, verbose=False
)
# Initialize the cache for storing variables needed for visulization.
visuals_cache = collections.deque(maxlen=cfg.visualization.history_size)
# Init the visualization visdom env.
if cfg.visualization.visdom:
viz = Visdom(
server=cfg.visualization.visdom_server,
port=cfg.visualization.visdom_port,
use_incoming_socket=False,
)
else:
viz = None
for epoch in range(cfg.optimizer.max_epochs):
image_translator.train()
stats.new_epoch()
for iteration,data in enumerate(train_dataloader):
optimizer.zero_grad()
views,param_vectors = data
views = views.float().cuda()
param_vectors = param_vectors.float().cuda()
camera_instance = Camera()
camera_instance.lookAt(param_vectors[0][0],math.degrees(param_vectors[0][1]),math.degrees(param_vectors[0][2]))
camera_location = param_vectors[0,3:6]
light_location = param_vectors[0,6:9]
torch_camera_face_attrs = torch_coord_face_attrs - camera_location
torch_light_face_attrs = torch_coord_face_attrs - light_location
torch_face_attrs = torch.cat([torch_camera_face_attrs,torch_normal_face_attrs,torch_light_face_attrs,torch_random_face_attrs],2)
rasterizer_instance = Rasterizer()
rasterizer_instance.init_rasterizer(camera_instance.camera)
fragments = rasterizer_instance.rasterizer(pytorch_mesh)
pix_to_face = fragments.pix_to_face
bary_coords = fragments.bary_coords
pix_features = torch.squeeze(interpolate_face_attributes(pix_to_face,bary_coords,torch_face_attrs),3)
predicted_render = image_translator(pix_features,torch_texture)
loss = 1000*mse_loss(predicted_render,views)
loss.backward()
optimizer.step()
# Update stats with the current metrics.
stats.update(
{"mse_loss": float(loss)},
stat_set="train",
)
if iteration % cfg.stats_print_interval == 0:
stats.print(stat_set="train")
# Adjust the learning rate.
#lr_scheduler.step()
# Validation
if epoch % cfg.validation_epoch_interval == 0: # and epoch > 0:
# Sample a validation camera/image.
val_batch = next(validation_dataloader.__iter__())
views, param_vectors= val_batch
views = views.float().cuda()
param_vectors = param_vectors.float().cuda()
# Activate eval mode of the model (allows to do a full rendering pass).
image_translator.eval()
with torch.no_grad():
camera_instance = Camera()
camera_instance.lookAt(param_vectors[0][0], math.degrees(param_vectors[0][1]), math.degrees(param_vectors[0][2]))
camera_location = param_vectors[0,3:6]
light_location = param_vectors[0,6:9]
torch_camera_face_attrs = torch_coord_face_attrs - camera_location
torch_light_face_attrs = torch_coord_face_attrs - light_location
torch_face_attrs = torch.cat([torch_camera_face_attrs,torch_normal_face_attrs,torch_light_face_attrs,torch_random_face_attrs],2)
rasterizer_instance = Rasterizer()
rasterizer_instance.init_rasterizer(camera_instance.camera)
fragments = rasterizer_instance.rasterizer(pytorch_mesh)
pix_to_face = fragments.pix_to_face
bary_coords = fragments.bary_coords
pix_features = torch.squeeze(interpolate_face_attributes(pix_to_face, bary_coords, torch_face_attrs), 3)
#pix_features = pix_features.permute(0, 3, 1, 2)
predicted_render = image_translator(pix_features,torch_texture)
loss = 1000*mse_loss(predicted_render,views)
# Update stats with the validation metrics.
stats.update({"mse_loss":loss}, stat_set="val")
stats.print(stat_set="val")
if viz is not None:
# Plot that loss curves into visdom.
stats.plot_stats(
viz=viz,
visdom_env=cfg.visualization.visdom_env,
plot_file=None,
)
# Visualize the intermediate results.
render_max = torch.max(predicted_render)
visualize_image_outputs(
validation_images = [views[0].permute(2,0,1),predicted_render[0].permute(2,0,1)],viz=viz,visdom_env=cfg.visualization.visdom_env
)
# Set the model back to train mode.
image_translator.train()
# Checkpoint.
if (
epoch % cfg.checkpoint_epoch_interval == 0
and len(cfg.checkpoint_path) > 0
and epoch > 0
):
print(f"Storing checkpoint {checkpoint_path}.")
data_to_store = {
"model": image_translator.state_dict(),
"features" : torch_face_attrs,
"optimizer": optimizer.state_dict(),
"stats": pickle.dumps(stats),
}
torch.save(data_to_store, checkpoint_path)
