python
复制代码
#!/usr/bin/env python
# -*- coding:utf-8 -*-
# @FileName :torch_learn2.py
# @Time :2024/11/16 19:53
# @Author :Jason Zhang
import torch
torch.manual_seed(12124)
if __name__ == "__main__":
run_code = 0
a = torch.rand([3, 2])
a_chunk1, a_chunk2 = torch.chunk(a, chunks=2)
a_chunk11, a_chunk21 = torch.chunk(a, chunks=2, dim=1)
print(f"a=\n{a}")
print(f"a_chunk1=\n{a_chunk1}")
print(f"a_chunk2=\n{a_chunk2}")
print(f"a_chunk11=\n{a_chunk11}")
print(f"a_chunk21=\n{a_chunk21}")
t = torch.tensor([[1, 2], [3, 4]])
t_gather = torch.gather(t, 1, torch.tensor([[0, 1], [1, 0]]))
print(f"t=\n{t}")
print(f"t_gather=\n{t_gather}")
reshape_12 = torch.arange(12).reshape((3, 4))
print(f"reshape_12=\n{reshape_12}")
reshape_11 = reshape_12.reshape((-1, 1))
print(f"reshape_11=\n{reshape_11}")
src = torch.arange(1, 11).reshape((2, 5))
index = torch.tensor([[0, 1, 2, 0]])
y = torch.zeros(3, 5, dtype=src.dtype).scatter_(0, index, src)
print(f"src=\n{src}")
print(f"y=\n{y}")
stack_a = torch.rand((3, 4))
stack_b = torch.rand((3, 4))
stack_ab = torch.stack((stack_a, stack_b))
print(f"stack_a=\n{stack_a}")
print(f"stack_b=\n{stack_b}")
print(f"stack_ab=\n{stack_ab},shape={stack_ab.shape}")
squeeze_1 = torch.rand((2, 1, 3))
squeeze_2 = torch.squeeze(squeeze_1)
print(f"squeeze_1.shape={squeeze_1.shape}")
print(f"squeeze_2.shape={squeeze_2.shape}")
python
复制代码
a=
tensor([[0.5555, 0.0484],
[0.3199, 0.2577],
[0.8874, 0.6888]])
a_chunk1=
tensor([[0.5555, 0.0484],
[0.3199, 0.2577]])
a_chunk2=
tensor([[0.8874, 0.6888]])
a_chunk11=
tensor([[0.5555],
[0.3199],
[0.8874]])
a_chunk21=
tensor([[0.0484],
[0.2577],
[0.6888]])
t=
tensor([[1, 2],
[3, 4]])
t_gather=
tensor([[1, 2],
[4, 3]])
reshape_12=
tensor([[ 0, 1, 2, 3],
[ 4, 5, 6, 7],
[ 8, 9, 10, 11]])
reshape_11=
tensor([[ 0],
[ 1],
[ 2],
[ 3],
[ 4],
[ 5],
[ 6],
[ 7],
[ 8],
[ 9],
[10],
[11]])
src=
tensor([[ 1, 2, 3, 4, 5],
[ 6, 7, 8, 9, 10]])
y=
tensor([[1, 0, 0, 4, 0],
[0, 2, 0, 0, 0],
[0, 0, 3, 0, 0]])
stack_a=
tensor([[0.2410, 0.9222, 0.5832, 0.3587],
[0.9344, 0.3320, 0.3852, 0.3239],
[0.7664, 0.9575, 0.2645, 0.5601]])
stack_b=
tensor([[0.4304, 0.7509, 0.3536, 0.7229],
[0.9026, 0.0793, 0.3076, 0.3272],
[0.4434, 0.2406, 0.7080, 0.9304]])
stack_ab=
tensor([[[0.2410, 0.9222, 0.5832, 0.3587],
[0.9344, 0.3320, 0.3852, 0.3239],
[0.7664, 0.9575, 0.2645, 0.5601]],
[[0.4304, 0.7509, 0.3536, 0.7229],
[0.9026, 0.0793, 0.3076, 0.3272],
[0.4434, 0.2406, 0.7080, 0.9304]]]),shape=torch.Size([2, 3, 4])
squeeze_1.shape=torch.Size([2, 1, 3])
squeeze_2.shape=torch.Size([2, 3])
Process finished with exit code 0