comparison of camera parameter optimization example of NMR and SoftRas. When use both model at same time, it will lead nan loss value.

camera_optimization.py

import os
import argparse
import glob

import torch
import torch.nn as nn
import numpy as np
from skimage.io import imread, imsave
import tqdm
import imageio

import neural_renderer as nr
import soft_renderer as sr

current_dir = os.path.dirname(os.path.realpath(__file__))
data_dir = os.path.join(current_dir, 'data')

global only_softras
only_softras = False

class softrasModel(nn.Module):
	def __init__(self, filename_obj, filename_ref=None):
		super(softrasModel, self).__init__()
		# load .obj
		vertices, faces = nr.load_obj(filename_obj)
		#self.register_buffer('vertices', vertices[None, :, :])
		#self.register_buffer('faces', faces[None, :, :])
		self.vertices = vertices[None, :, :].cuda()
		self.faces = faces[None, :, :].cuda()

		# create textures
		textures = 	torch.ones(1, 2464, 1, 3, dtype=torch.float32) # [1, 2700, 1, 3]
		#self.register_buffer('textures', textures)
		self.textures = textures.cuda()

		# load reference image
		image_ref = torch.from_numpy((imread(filename_ref).max(-1) != 0).astype(np.float32))
		#self.register_buffer('image_ref', image_ref)
		self.image_ref = image_ref.cuda()

		# camera parameters
		self.camera_position = nn.Parameter(torch.from_numpy(np.array([6, 10, -14], dtype=np.float32)))
		print("softras nn.Parameter ID: ", id(self.camera_position))

		# setup renderer
		sigma = 1e-5
		gamma = 1e-4
		renderer = sr.SoftRenderer(camera_mode='look_at', perspective=True, eye=self.camera_position, sigma_val=sigma, gamma_val=gamma)
		self.renderer = renderer

	def forward(self):
		image = self.renderer(self.vertices, self.faces)[:,-1,:,:]
		# print(image.shape)  # torch.Size([1, 256, 256])

		loss = torch.sum((image - self.image_ref[None, :, :]) ** 2)
		return loss

class nmrModel(nn.Module):
	def __init__(self, filename_obj, filename_ref=None):
		super(nmrModel, self).__init__()
		# load .obj
		vertices, faces = nr.load_obj(filename_obj)
		# self.register_buffer('vertices', vertices[None, :, :])
		# self.register_buffer('faces', faces[None, :, :])
		self.vertices = vertices[None, :, :].cuda()
		self.faces = faces[None, :, :].cuda()

		# create textures
		texture_size = 2
		textures = torch.ones(1, self.faces.shape[1], texture_size, texture_size, texture_size, 3, dtype=torch.float32)
		
		# self.register_buffer('textures', textures)
		self.textures = textures.cuda()

		# load reference image
		image_ref = torch.from_numpy((imread(filename_ref).max(-1) != 0).astype(np.float32))
		# self.register_buffer('image_ref', image_ref)
		self.image_ref = image_ref.cuda()

		# camera parameters
		self.camera_position = nn.Parameter(torch.from_numpy(np.array([6, 10, -14], dtype=np.float32)))
		print("nmr nn.Parameter ID: ", id(self.camera_position))

		# setup renderer
		renderer = nr.Renderer(camera_mode='look_at', perspective=True)
		renderer.eye = self.camera_position
		
		self.renderer = renderer

	def forward(self):
		image = self.renderer(self.vertices, self.faces, mode='silhouettes')
		# print(image.shape) # torch.Size([1, 256, 256])

		loss = torch.sum((image - self.image_ref[None, :, :]) ** 2)
		return loss

def make_gif(filename):
	with imageio.get_writer(filename, mode='I') as writer:
		for filename in sorted(glob.glob('/tmp/_tmp_*.png')):
			writer.append_data(imread(filename))
			os.remove(filename)
	writer.close()


# def make_reference_image(filename_ref, filename_obj, use_nmr):
# 	model = Model(filename_obj)
# 	model.cuda()

# 	model.renderer.eye = nr.get_points_from_angles(2.732, 30, -15)
# 	images = model.renderer.render(model.vertices, model.faces, torch.tanh(model.textures))
# 	image = images.detach().cpu().numpy()[0]
# 	imsave(filename_ref, image)

# def clip_grad_value_(parameters, clip_value):
# 	if isinstance(parameters, torch.Tensor):
# 		parameters = [parameters]
# 	clip_value = float(clip_value)
# 	for p in filter(lambda p: p.grad is not None, parameters):
# 		for i in range(p.size(0)): 
# 			if (0 <= p.grad.data[i]) and (p.grad.data[i] < clip_value):
# 				p.grad.data[i].clamp_(clip_value=+clip_value)
# 			elif (-clip_value < p.grad.data[i]) and (p.grad.data[i] < 0):
# 				p.grad.data[i].clamp_(clip_value=-clip_value)
# 			else:
# 				pass

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)

	if not only_softras:
		nmr_model = nmrModel(args.filename_obj, args.filename_ref)
		nmr_model.cuda()

	sr_model = softrasModel(args.filename_obj, args.filename_ref)
	sr_model.cuda()

	# optimizer = chainer.optimizers.Adam(alpha=0.1)
	if not only_softras: 
		nmr_optimizer = torch.optim.Adam(nmr_model.parameters(), lr=0.1)
	sr_optimizer = torch.optim.Adam(sr_model.parameters(), lr=0.1)

	loop = tqdm.tqdm(range(300))
	for i in loop:
		if not only_softras: 
			nmr_optimizer.zero_grad()
		sr_optimizer.zero_grad()

		if not only_softras: 
			nmr_loss = nmr_model()
		sr_loss = sr_model()

		
		if not only_softras: 
			nmr_loss.backward()
		sr_loss.backward()

		# minimal clip
		#clip_grad_value_(nmr_model.parameters(), 1e-4)
		#clip_grad_value_(sr_model.parameters(), 1e-4)

		# max clip
		#max_norm = 5.0
		#torch.nn.utils.clip_grad_norm_(nmr_model.parameters(), max_norm, norm_type=2)
		#torch.nn.utils.clip_grad_norm_(sr_model.parameters(), max_norm, norm_type=2)

		if not only_softras: 
			nmr_optimizer.step()
		sr_optimizer.step()

		if not only_softras: 
			nmr_images = nmr_model.renderer(nmr_model.vertices, nmr_model.faces, torch.tanh(nmr_model.textures))
		sr_images = sr_model.renderer(sr_model.vertices, sr_model.faces, torch.tanh(sr_model.textures))
		
		nmr_image = nmr_images[0][0,:,:,:].detach().cpu() if not only_softras else torch.zeros_like(sr_images)[0,0:3,:,:].detach().cpu()
		sr_image = sr_images[0,0:3,:,:].detach().cpu()

		image = torch.cat((nmr_image, sr_image), dim=-1).numpy().transpose(1,2,0)
		imsave('/tmp/_tmp_%04d.png' % i, image)
		if not only_softras:
			loop.set_description('Optimizing (nmr_loss {:0.4f}, sr_loss {:0.4f})'.format(nmr_loss.data, sr_loss.data))
			#if (nmr_loss.item() < 70) and (sr_loss.item() < 70): break
		else:
			loop.set_description('Optimizing (sr_loss {:0.4f})'.format(sr_loss.data))
			#if sr_loss.item() < 70: break
		
	make_gif(args.filename_output)


if __name__ == '__main__':
  main()