# 使用MindSpore替换torch.distributions的Categorical函数

**1. 系统环境**

硬件环境\(Ascend/GPU/CPU\): Ascend/GPU/CPU

软件环境: MindSpore 版本: 2.0.0

执行模式（静态图/动态图）： 不限

Python 版本: 3.7.6

操作系统平台: linux

**2. 报错信息**

2.1 问题描述

将pytorch转换成mindspore代码的时候，调用了pytorch.distributions的Categorical方法，这三行代码该怎么使用mindspore的方法进行呢。

2.3 脚本信息

```
import torch.nn.functional as F
from torch.distributions import Categorical
def sample_cycle_component(self, logits, output_type, random_ratio=0):
        action_index_local = Categorical(logits=logits).sample()
        prob_matrix = F.softmax(logits, dim=1)
        log_prob_matrix = F.log_softmax(logits, dim=1)
```

**3. 根因分析**

[torch.distributions.Categorical](https://pytorch.org/docs/stable/distributions.html#torch.distributions.categorical.Categorical)  =\>  [mindspore.nn.probability.distribution.Categorical](https://www.mindspore.cn/docs/zh-CN/r2.0/api_python/nn_probability/mindspore.nn.probability.distribution.Categorical.html)

[torch.nn.functional.softmax](https://pytorch.org/docs/stable/generated/torch.nn.functional.softmax.html#torch.nn.functional.softmax)  =\>  [mindspore.ops.softmax](https://www.mindspore.cn/docs/zh-CN/r2.0/api_python/ops/mindspore.ops.softmax.html)

[torch.nn.functional.log\_softmax ](https://pytorch.org/docs/stable/generated/torch.nn.functional.log_softmax.html#torch.nn.functional.log_softmax)=\>  [mindspore.ops.log\_softmax](https://www.mindspore.cn/docs/zh-CN/r2.0/api_python/ops/mindspore.ops.log_softmax.html)

**4. 解决方案**

```
from mindspore import ops
from mindspore.nn.probability.distribution import Categorical
def sample_cycle_component(logits, output_type, random_ratio=0):
    action_index_local = Categorical(probs=input).sample()
    prob_matrix = ops.softmax(logits, axis=1)
    log_prob_matrix = ops.log_softmax(logits, axis=1)
```

