More dedicated details to optimize speed

[lucasjinreal]

Hi, as I am testing, panopatic-deeplab can get a satisfying speed in terms of model forward time. However, when counts the postprocessing and visualize speed (this part can not be ignored since it should not be so slow), the speed is totally slow.

Wanna ask, if you guys have a more dedicated optimized plan / method, would be nice to optimize this model and so that more promising converting to other framework for deployment such as accelerate it with TensorRT and deploy more lightweighted model.

[bowenc0221]

The for-loops in this function could be optimized to further speed-up post-processing time.

panoptic-deeplab/segmentation/model/post_processing/instance_post_processing.py

Lines 123 to 179 in 5b3dd8c

	def merge_semantic_and_instance(sem_seg, ins_seg, label_divisor, thing_list, stuff_area, void_label):
	"""
	Post-processing for panoptic segmentation, by merging semantic segmentation label and class agnostic
	instance segmentation label.
	Arguments:
	sem_seg: A Tensor of shape [1, H, W], predicted semantic label.
	ins_seg: A Tensor of shape [1, H, W], predicted instance label.
	label_divisor: An Integer, used to convert panoptic id = semantic id * label_divisor + instance_id.
	thing_list: A List of thing class id.
	stuff_area: An Integer, remove stuff whose area is less tan stuff_area.
	void_label: An Integer, indicates the region has no confident prediction.
	Returns:
	A Tensor of shape [1, H, W] (to be gathered by distributed data parallel).
	Raises:
	ValueError, if batch size is not 1.
	"""
	# In case thing mask does not align with semantic prediction
	pan_seg = torch.zeros_like(sem_seg) + void_label
	thing_seg = ins_seg > 0
	semantic_thing_seg = torch.zeros_like(sem_seg)
	for thing_class in thing_list:
	semantic_thing_seg[sem_seg == thing_class] = 1
	# keep track of instance id for each class
	class_id_tracker = {}
	# paste thing by majority voting
	instance_ids = torch.unique(ins_seg)
	for ins_id in instance_ids:
	if ins_id == 0:
	continue
	# Make sure only do majority voting within semantic_thing_seg
	thing_mask = (ins_seg == ins_id) & (semantic_thing_seg == 1)
	if torch.nonzero(thing_mask).size(0) == 0:
	continue
	class_id, _ = torch.mode(sem_seg[thing_mask].view(-1, ))
	if class_id.item() in class_id_tracker:
	new_ins_id = class_id_tracker[class_id.item()]
	else:
	class_id_tracker[class_id.item()] = 1
	new_ins_id = 1
	class_id_tracker[class_id.item()] += 1
	pan_seg[thing_mask] = class_id * label_divisor + new_ins_id
	# paste stuff to unoccupied area
	class_ids = torch.unique(sem_seg)
	for class_id in class_ids:
	if class_id.item() in thing_list:
	# thing class
	continue
	# calculate stuff area
	stuff_mask = (sem_seg == class_id) & (~thing_seg)
	area = torch.nonzero(stuff_mask).size(0)
	if area >= stuff_area:
	pan_seg[stuff_mask] = class_id * label_divisor
	return pan_seg