Differences with torch.nn.functional.softmax

View Source On Gitee

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)