Computer >> Hướng Dẫn Máy Tính >  >> Lập Trình >> Python

Cách tìm phần tử thứ k và top k phần tử của tensor trong PyTorch?

PyTorch cung cấp phương thức torch.kthvalue() để tìm phần tử thứ k của tensor. Nó trả về giá trị phần tử thứ k của tensor đã sắp xếp tăng dần, và chỉ mục của phần tử trong tensor gốc.

Phương thức torch.topk() dùng để tìm top "k" phần tử. Nó trả về k phần tử lớn nhất trong tensor.

Các bước

  • Nhập thư viện torch. Đảm bảo đã cài.

  • Tạo tensor PyTorch và in ra.

  • Tính torch.kthvalue(input, k). Trả về hai tensor. Gán cho biến value và index. input là tensor, k là số nguyên.

  • Tính torch.topk(input, k). Trả về hai tensor: giá trị top k và chỉ mục của chúng trong tensor gốc. Gán cho values và indices.

  • In giá trị và chỉ mục phần tử thứ k, và giá trị, chỉ mục top k.

Ví dụ 1

Chương trình tìm phần tử thứ k của tensor.

import torch

# Tạo tensor 1D
T = torch.Tensor([2.334, 4.433, -4.33, -0.433, 5, 4.443])
print("Original Tensor:\n", T)
# Tìm phần tử thứ 3 trong tensor đã sắp xếp
value, index = torch.kthvalue(T, 3)
print("3rd element value:", value)
print("3rd element index:", index)

Kết quả

Original Tensor:
tensor([ 2.3340,  4.4330, -4.3300, -0.4330,  5.0000,  4.4430])
3rd element value: tensor(2.3340)
3rd element index: tensor(0)

Ví dụ 2

Chương trình tìm top k phần tử lớn nhất.

import torch

T = torch.Tensor([2.334, 4.433, -4.33, -0.433, 5, 4.443])
print("Original Tensor:\n", T)
# Tìm top k=2 phần tử lớn nhất
values, indices = torch.topk(T, 2)
print("Top 2 element values:", values)
print("Top 2 element indices:", indices)

Kết quả

Original Tensor:
tensor([ 2.3340,  4.4330, -4.3300, -0.4330,  5.0000,  4.4430])
Top 2 element values: tensor([5.0000, 4.4430])
Top 2 element indices: tensor([4, 5])