Scatter is potentially faster than one-hot

[ATATC]

onehot = torch.zeros(output.shape, device=output.device, dtype=torch.float16)
onehot.scatter_(1, output_seg, 1)

The above implementation is very likely to be faster than the current implementation of convert_ids_to_logits().

[ATATC]

@perctrix Can you verify?

[perctrix]

I've verified that: scatter_ is indeed significantly faster than one_hot. Here's the result table:

Config	Shape	Classes	CPU one_hot	CPU scatter	CPU Speedup	CUDA one_hot	CUDA scatter	CUDA Speedup
2D small	(4, 256, 256)	2	1.379ms	0.405ms	3.40x	0.252ms	0.116ms	2.17x
2D medium	(4, 512, 512)	4	8.539ms	1.728ms	4.94x	0.815ms	0.093ms	8.76x
2D large	(8, 512, 512)	8	44.605ms	5.689ms	7.84x	3.321ms	0.744ms	4.46x
3D small	(2, 64, 64, 64)	2	2.588ms	0.792ms	3.27x	0.303ms	0.126ms	2.41x
3D medium	(2, 128, 128, 128)	4	33.210ms	8.765ms	3.79x	5.314ms	0.936ms	5.67x
3D large	(1, 256, 256, 256)	8	313.856ms	53.722ms	5.84x	38.577ms	5.968ms	6.46x

Reason: one_hot requires movedim + contiguous() to rearrange the memory after generating the result for the last dimension, while scatter_ writes directly in-place on the correct memory layout, saving extra memory allocation and copying. The larger the tensor and the more classes, the more obvious the advantage.

[perctrix]

@ATATC