torch.unsqueeze设定维度值
时间: 2023-10-16 16:00:21 浏览: 85
浅谈pytorch中torch.max和F.softmax函数的维度解释
torch.unsqueeze函数可以用于在指定维度上增加一个维度,使得原来的张量的维度增加1。它的语法是:
```python
torch.unsqueeze(input, dim)
```
其中,`input`是输入的张量,`dim`是要在其上增加维度的维度索引。
举个例子,假设我们有一个形状为(3, 4)的二维张量`x`:
```python
x = torch.tensor([[1, 2, 3, 4],
[5, 6, 7, 8],
[9, 10, 11, 12]])
```
如果我们想在第0维度上增加一个维度,可以使用`torch.unsqueeze(x, 0)`。这将返回一个形状为(1, 3, 4)的三维张量,其中第0个维度的大小为1。
同样地,如果我们想在第1维度上增加一个维度,可以使用`torch.unsqueeze(x, 1)`。这将返回一个形状为(3, 1, 4)的三维张量,其中第1个维度的大小为1。
需要注意的是,`dim`参数的取值范围应该在[-input.dim()-1, input.dim()+1]之间,否则会抛出索引越界的异常。
阅读全文