Fused Jasper block

[fedorovgv]

Hi everyone! In Jasper paper i found this lines:

Each sub-block can be fused into a single GPU kernel: dropout is not used at inference-time and is
eliminated, batch norm can be fused with the preceding convolution, ReLU clamps the result, and
residual summation can be treated as a modified bias term in this fused operation.

But in the jasper block code in nemo, I didn't see this sub-block merging completely into the core. Do I understand correctly that the line "Each sub-block can be fused" means that the operations in the sub-block are fused separately (conv->batch_norm->relu->dropout and conv fused, then the results are returned in batch_norm, which are also fused, etc.), not the entire sub-block?

No attempt is made to completely fused the subblock.

NeMo/nemo/collections/asr/parts/submodules/jasper.py

Lines 974 to 1051 in 5d5c0fb

	def forward(self, input_: Tuple[List[Tensor], Optional[Tensor]]):
	"""
	Forward pass of the module.
	Args:
	input_: The input is a tuple of two values - the preprocessed audio signal as well as the lengths
	of the audio signal. The audio signal is padded to the shape [B, D, T] and the lengths are
	a torch vector of length B.
	Returns:
	The output of the block after processing the input through `repeat` number of sub-blocks,
	as well as the lengths of the encoded audio after padding/striding.
	"""
	# type: (Tuple[List[Tensor], Optional[Tensor]]) -> Tuple[List[Tensor], Optional[Tensor]] # nopep8
	lens_orig = None
	xs = input_[0]
	if len(input_) == 2:
	xs, lens_orig = input_
	# compute forward convolutions
	out = xs[-1]
	lens = lens_orig
	for i, l in enumerate(self.mconv):
	# if we're doing masked convolutions, we need to pass in and
	# possibly update the sequence lengths
	# if (i % 4) == 0 and self.conv_mask:
	if isinstance(l, (MaskedConv1d, SqueezeExcite)):
	out, lens = l(out, lens)
	else:
	out = l(out)
	# compute the residuals
	if self.res is not None:
	for i, layer in enumerate(self.res):
	res_out = xs[i]
	for j, res_layer in enumerate(layer):
	if isinstance(res_layer, MaskedConv1d):
	res_out, _ = res_layer(res_out, lens_orig)
	else:
	res_out = res_layer(res_out)
	if self.residual_mode == 'add' or self.residual_mode == 'stride_add':
	if PYTORCH_QUANTIZATION_AVAILABLE and self.quantize:
	out = self.residual_quantizer(out) + res_out
	elif not PYTORCH_QUANTIZATION_AVAILABLE and self.quantize:
	raise ImportError(
	"pytorch-quantization is not installed. Install from "
	"https://github.com/NVIDIA/TensorRT/tree/master/tools/pytorch-quantization."
	)
	else:
	out = out + res_out
	else:
	out = torch.max(out, res_out)
	# compute the output
	out = self.mout(out)
	# Support ASR Adapters
	if self.is_adapter_available():
	# Check for all available and enabled adapters
	adapter_names = self.get_enabled_adapters()
	if len(adapter_names) > 0:
	out = out.transpose(1, 2) # (B, T, C)
	# Call the adapters
	out = self.forward_enabled_adapters(out)
	out = out.transpose(1, 2) # (B, C, T)
	if self.is_access_enabled() and self.access_cfg.get('save_encoder_tensors', False):
	self.register_accessible_tensor(name='encoder', tensor=out)
	if self.res is not None and self.dense_residual:
	return xs + [out], lens
	return [out], lens

Can someone help me. I'm thinking, maybe these are some other details responsible for the fusion of the entire sub-block? or does "fused" mean that the subblock consists of "fused" operations like conv and batch norm?

[titu1994]

The paper mentions it as a possibility, say if you trace it and optimize the operations it would be possible to fuse it. However in pytorch at that time, there was no such jit tracer. Nowadays you can do jit trace of pytorch modules and while it will not fuse all the operations, it is feasible that in the future onnx or other optimizing jit compilers can do it automatically v

[fedorovgv]

Thanks a lot for the quick and complete answer!