Differences with torch.nn.functional.softmax

torch.nn.functional.softmax

torch.nn.functional.softmax(
    input,
    dim=None,
    _stacklevel=3,
    dtype=None
)

For more information, see torch.nn.functional.softmax.

mindspore.ops.softmax

class mindspore.ops.softmax(
    x,
    axis=-1
)

For more information, see mindspore.ops.softmax.

Differences

PyTorch: Supports to implement the function with the dim parameter and input, scaling the specified dimension elements between [0, 1] and the total to 1.

MindSpore: Supports to implement the function with the axis parameter and x, scaling the specified dimension elements between [0, 1] and the total to 1.

Classification	Subclass	PyTorch	MindSpore	difference
parameter	parameter1	input	x	Same functions, different parameter names
	parameter2	dim	axis	Same functions, different parameter names
	parameter3	dtype	-	Uesd to specify the type of data for the output Tensor in PyTorch. This parameter does not exist in MindSpore

Code Example

import mindspore as ms
import mindspore.ops as ops
import torch
import numpy as np

# In MindSpore, we can define an instance of this class first, and the default value of the parameter axis is -1.
x = ms.Tensor(np.array([1, 2, 3, 4, 5]), ms.float32)
output1 = ops.softmax(x)
print(output1)
# Out:
# [0.01165623 0.03168492 0.08612853 0.23412165 0.63640857]
x = ms.Tensor(np.array([[1, 2, 3, 4, 5], [5, 4, 3, 2, 1]]), ms.float32)
output2 = ops.softmax(x, axis=0)
print(output2)
# Out:
# [[0.01798621 0.11920292 0.5        0.880797   0.98201376]
#  [0.98201376 0.880797   0.5        0.11920292 0.01798621]]

# In torch, the input and dim should be input at the same time to implement the function.
input = torch.tensor(np.array([1.0, 2.0, 3.0, 4.0, 5.0]))
output3 = torch.nn.functional.softmax(input, dim=0)
print(output3)
# Out:
# tensor([0.0117, 0.0317, 0.0861, 0.2341, 0.6364], dtype=torch.float64)