代码:
import torch
input = torch.randn(2,2,3,3) # 生成随机张量
print("input",input)
print("input.size",input.size())
output = torch.nn.Softmax(dim=1)(input) # 沿着维度1进行Softmax计算
print("output", output)
print("output.size", output.size())
结果:
input tensor([[[[-1.0284, 0.3340, 0.8394],
[ 2.0868, -0.4440, 1.1711],
[-0.4021, -1.5332, 1.3740]],
[[-1.7498, -2.3110, -0.9069],
[ 1.3210, -2.1890, 0.8349],
[-0.0612, 0.3885, -1.3221]]],
[[[ 0.9202, 0.4145, 0.5510],
[-1.2614, 1.1752, -0.7427],
[ 0.1773, -0.1683, -0.4466]],
[[ 1.7270, 0.7876, -0.0378],
[ 1.1749, 0.1121, -0.0742],
[-1.4793, -0.5104, 1.9719]]]])
input.size torch.Size([2, 2, 3, 3])
output tensor([[[[0.6729, 0.9337, 0.8515],
[0.6826, 0.8513, 0.5833],
[0.4156, 0.1277, 0.9368]],
[[0.3271, 0.0663, 0.1485],
[0.3174, 0.1487, 0.4167],
[0.5844, 0.8723, 0.0632]]],
[[[0.3086, 0.4078, 0.6431],
[0.0804, 0.7433, 0.3388],
[0.8398, 0.5847, 0.0818]],
[[0.6914, 0.5922, 0.3569],
[0.9196, 0.2567, 0.6612],
[0.1602, 0.4153, 0.9182]]]])
output.size torch.Size([2, 2, 3, 3]) # 维度1之和为1