import torch

a = torch.tensor([[1, 2], [3, 4]])  # 形狀 (2, 2)
b = torch.tensor([[5, 6], [7, 8]])

# 在第0維拼接（垂直方向）
c = torch.cat([a, b], dim=0)  
print(c)
# 輸出:
# tensor([[1, 2],
#         [3, 4],
#         [5, 6],
#         [7, 8]])

# 在第1維拼接（水平方向）
d = torch.cat([a, b], dim=1)  
print(d)
# 輸出:
# tensor([[1, 2, 5, 6],
#         [3, 4, 7, 8]])

a = torch.tensor([1, 2, 3])
b = torch.tensor([4, 5, 6])

# 沿第0維堆疊，生成二維張量
c = torch.stack([a, b], dim=0)  
print(c.shape)  # torch.Size([2, 3])
print(c)
# 輸出:
# tensor([[1, 2, 3],
#         [4, 5, 6]])

# 沿第1維堆疊，生成二維張量
d = torch.stack([a, b], dim=1)  
print(d.shape)  # torch.Size([3, 2])
print(d)
# 輸出:
# tensor([[1, 4],
#         [2, 5],
#         [3, 6]])

a = torch.arange(9)  # tensor([0, 1, 2, 3, 4, 5, 6, 7, 8])

# 按列表尺寸分割 [2,3,4]
parts = torch.split(a, [2, 3, 4], dim=0)
for part in parts:
    print(part)

'''
輸出:
tensor([0, 1])
tensor([2, 3, 4])
tensor([5, 6, 7, 8])
'''

# 平均分割為3份
chunks = torch.split(a, 3, dim=0)
print([c.shape for c in chunks])  # [torch.Size([3]), torch.Size([3]), torch.Size([3])]

a = torch.arange(10)  # tensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])

# 分成3份，默認(rèn)在第0維操作
chunks = torch.chunk(a, chunks=3, dim=0)
for i, chunk in enumerate(chunks):
    print(f"Chunk {i}: {chunk}")

'''
輸出:
Chunk 0: tensor([0, 1, 2, 3])
Chunk 1: tensor([4, 5, 6])
Chunk 2: tensor([7, 8, 9])
'''

# 在第1維分割二維張量
b = a.reshape(2,5)
chunks = torch.chunk(b, chunks=2, dim=1)
print(chunks[0].shape)  # torch.Size([2, 2])
print(chunks[1].shape)  # torch.Size([2, 3])

import torch
from torchvision import transforms
from PIL import Image
import matplotlib.pyplot as plt

# 1. 加載示例圖像 (H, W, C) -> 轉(zhuǎn)換為 (C, H, W)
image = Image.open('cat.jpg').convert('RGB')
image = transforms.ToTensor()(image)  # shape: torch.Size([3, 256, 256])

# 2. 使用split分離RGB通道
r_channel, g_channel, b_channel = torch.split(image, split_size_or_sections=1, dim=0)

''' 可視化原始通道
plt.figure(figsize=(12,4))
plt.subplot(131), plt.imshow(r_channel.squeeze().numpy(), cmap='Reds'), plt.title('Red')
plt.subplot(132), plt.imshow(g_channel.squeeze().numpy(), cmap='Greens'), plt.title('Green')
plt.subplot(133), plt.imshow(b_channel.squeeze().numpy(), cmap='Blues'), plt.title('Blue')
plt.show()
'''

# 3. 對每個通道進(jìn)行歸一化（示例操作）
def normalize(tensor):
    return (tensor - tensor.mean()) / tensor.std()

r_norm = normalize(r_channel)
g_norm = normalize(g_channel)
b_norm = normalize(b_channel)

# 4. 使用cat合并處理后的通道
normalized_img = torch.cat([r_norm, g_norm, b_norm], dim=0)
'''觀察歸一化效果
plt.imshow(normalized_img.permute(1,2,0))
plt.title('Normalized Image')
plt.show()
'''

# 5. 創(chuàng)建模擬圖像批次 (假設(shè)有4張相同圖像)
batch_images = torch.stack([image]*4, dim=0)  # shape: (4, 3, 256, 256)

# 6. 使用chunk分割批次為訓(xùn)練集/驗證集
train_set, val_set = torch.chunk(batch_images, chunks=2, dim=0)
print(f"Train set size: {train_set.shape}")  # torch.Size([2, 3, 256, 256])
print(f"Val set size: {val_set.shape}")      # torch.Size([2, 3, 256, 256])