PyTorch中的torch.cat函數基本用法詳解

更新時間：2024年08月19日 15:23:18 作者：小橋流水---人工智能

在PyTorch中,torch.cat是一個非常實用的函數,用于將多個張量（Tensor）沿指定維度連接起來,本文將詳細介紹torch.cat函數的用法,并通過一些示例來說明其應用,感興趣的朋友跟隨小編一起看看吧

在PyTorch中，torch.cat是一個非常實用的函數，用于將多個張量（Tensor）沿指定維度連接起來。這個功能在機器學習和深度學習中經常用到，尤其是在需要合并數據或模型輸出時。本文將詳細介紹torch.cat函數的用法，并通過一些示例來說明其應用。

1. torch.cat的基本用法

torch.cat的基本語法如下：

torch.cat(tensors, dim=0, out=None)

tensors：一個張量序列，可以是任何形式的Python序列，如列表或元組。
dim：要連接的維度。在PyTorch中，每個維度都有一個索引，從0開始。
out：可選參數，用于指定輸出張量。

2. 示例

讓我們通過一些示例來看看如何使用torch.cat。

示例 1：連接一維張量

import torch
# 創(chuàng)建一維張量
a = torch.tensor([1, 2, 3])
b = torch.tensor([4, 5, 6])
# 沿著第0維連接
result = torch.cat((a, b), dim=0)
print(result)  # 輸出：tensor([1, 2, 3, 4, 5, 6])

這個例子中，兩個一維張量沿著第0維連接，結果就是將它們首尾相接。

示例 2：連接二維張量

# 創(chuàng)建二維張量
a = torch.tensor([[1, 2], [3, 4]])
b = torch.tensor([[5, 6], [7, 8]])
# 沿著第0維連接
result0 = torch.cat((a, b), dim=0)
print(result0)
# 輸出：
# tensor([[1, 2],
#         [3, 4],
#         [5, 6],
#         [7, 8]])
# 沿著第1維連接
result1 = torch.cat((a, b), dim=1)
print(result1)
# 輸出：
# tensor([[1, 2, 5, 6],
#         [3, 4, 7, 8]])

在這個示例中，兩個二維張量分別沿著第0維和第1維進行連接。沿著第0維連接就像是在垂直方向上疊加矩陣，而沿著第1維連接則是在水平方向上拼接它們。