torch.chunk的用法示例——非常清晰

`torch.chunk` 是 PyTorch 中用于将一个张量沿指定维度分割成多个子张量的函数。下面是对其用法的详细解释和示例。

函数定义

```python

torch.chunk(input, chunks, dim=0)

```

  • **`input`**: 要分割的输入张量。

  • **`chunks`**: 要分割成的子张量的数量。

  • **`dim`**: 指定在哪个维度上进行分割,默认为 `0`(第一个维度,即行)。

示例

示例 1: 基本用法

复制代码
import torch

# 创建一个 4x4 的张量
tensor = torch.arange(16).reshape(4, 4)
print("Original Tensor:")
print(tensor)

# 将张量分割成 2 个部分,沿第 0 维(行)
chunks = torch.chunk(tensor, 2, dim=0)

print("\nChunks:")
for i, chunk in enumerate(chunks):
    print(f"Chunk {i}:")
    print(chunk)

**输出**:

```

Original Tensor:

tensor([[ 0, 1, 2, 3],

4, 5, 6, 7\], \[ 8, 9, 10, 11\], \[12, 13, 14, 15\]\]) Chunks: Chunk 0: tensor(\[\[0, 1, 2, 3\], \[4, 5, 6, 7\]\]) Chunk 1: tensor(\[\[ 8, 9, 10, 11\], \[12, 13, 14, 15\]\]) \`\`\`

示例 2: 不同维度分割

复制代码
# 将张量分割成 4 个部分,沿第 1 维(列)
chunks = torch.chunk(tensor, 4, dim=1)

print("\nChunks along dim=1:")
for i, chunk in enumerate(chunks):
    print(f"Chunk {i}:")
    print(chunk)

**输出**:

```

Chunks along dim=1:

Chunk 0:

tensor([[0],

4\], \[8\], \[12\]\]) Chunk 1: tensor(\[\[ 1\], \[ 5\], \[ 9\], \[13\]\]) Chunk 2: tensor(\[\[ 2\], \[ 6\], \[10\], \[14\]\]) Chunk 3: tensor(\[\[ 3\], \[ 7\], \[11\], \[15\]\]) \`\`\`

总结

  • `torch.chunk` 可以方便地将张量按指定维度分割成多个子张量,适用于需要将数据划分为多个部分的情况。

  • 在处理深度学习任务时,这种分割操作可以帮助实现特定的特征处理或聚合策略。

相关推荐
我爱一条柴ya3 分钟前
【AI大模型】神经网络反向传播:核心原理与完整实现
人工智能·深度学习·神经网络·ai·ai编程
万米商云7 分钟前
企业物资集采平台解决方案:跨地域、多仓库、百部门——大型企业如何用一套系统管好百万级物资?
大数据·运维·人工智能
新加坡内哥谈技术10 分钟前
Google AI 刚刚开源 MCP 数据库工具箱,让 AI 代理安全高效地查询数据库
人工智能
慕婉030712 分钟前
深度学习概述
人工智能·深度学习
大模型真好玩13 分钟前
准确率飙升!GraphRAG如何利用知识图谱提升RAG答案质量(额外篇)——大规模文本数据下GraphRAG实战
人工智能·python·mcp
198914 分钟前
【零基础学AI】第30讲:生成对抗网络(GAN)实战 - 手写数字生成
人工智能·python·深度学习·神经网络·机器学习·生成对抗网络·近邻算法
6confim14 分钟前
AI原生软件工程师
人工智能·ai编程·cursor
阿里云大数据AI技术15 分钟前
Flink Forward Asia 2025 主旨演讲精彩回顾
大数据·人工智能·flink
i小溪16 分钟前
在使用 Docker 时,如果容器挂载的数据目录(如 `/var/moments`)位于数据盘,只要服务没有读写,数据盘是否就不会被唤醒?
人工智能·docker
程序员NEO18 分钟前
Spring AI 对话记忆大揭秘:服务器重启,聊天记录不再丢失!
人工智能·后端