PyTorch搜索Tensor指定维度的前K大个(K小个)元素--------(torch.topk)命令参数详解及举例

torch.topk

语法

torch.topk(input, k, dim=None, largest=True, sorted=True, *, out = None)

作用

返回输入tensorinput中,在给定的维度dimk个最大的元素。

如果dim没有给定,那么选择输入input的最后一维。

如果largest = False,那么返回k个最小的元素。

返回一个namedtuple类型的元组(values, indices),其中indices是指元素在原数组中的索引。

sorted = True, 则返回的k个元素是有序的。

Parameters

  • input (Tensor) – the input tensor
    输入的张量

  • k (int) – the k in “top-k”
    返回的k的值

  • dim(int, optional) – the dimension to sort along
    指定的排序的维度 ,如果dim没有给定,那么选择输入input的最后一维。 dim若为-1,文档未说明,但是根据实操效果,应该也是对最后一维进行search。
    如shape为Batch_size x p x q,返回结果为Batch_size x p x k

  • largest(bool, optional) – controls whether to return largest or smallest elements
    True返回最大值,False返回最小值。

  • sorted(bool, optional) – controls whether to return the elements in sorted order
    控制返回的元素是否排序。

例子

>>> x = torch.arange(1., 6.)
>>> x
tensor([ 1.,  2.,  3.,  4.,  5.])
>>> torch.topk(x, 3)
torch.return_types.topk(values=tensor([5., 4., 3.]), indices=tensor([4, 3, 2]))
  • 5
    点赞
  • 12
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
好的,以下是一个简单的encoder-decoder模型的代码,仅使用PyTorch和PyCharm: ``` import torch import torch.nn as nn import torch.optim as optim # 定义 Encoder 模型 class Encoder(nn.Module): def __init__(self, input_size, hidden_size): super(Encoder, self).__init__() self.hidden_size = hidden_size self.embedding = nn.Embedding(input_size, hidden_size) self.gru = nn.GRU(hidden_size, hidden_size) def forward(self, input, hidden): embedded = self.embedding(input).view(1, 1, -1) output = embedded output, hidden = self.gru(output, hidden) return output, hidden def init_hidden(self): return torch.zeros(1, 1, self.hidden_size) # 定义 Decoder 模型 class Decoder(nn.Module): def __init__(self, hidden_size, output_size): super(Decoder, self).__init__() self.hidden_size = hidden_size self.embedding = nn.Embedding(output_size, hidden_size) self.gru = nn.GRU(hidden_size, hidden_size) self.out = nn.Linear(hidden_size, output_size) self.softmax = nn.LogSoftmax(dim=1) def forward(self, input, hidden): output = self.embedding(input).view(1, 1, -1) output = nn.functional.relu(output) output, hidden = self.gru(output, hidden) output = self.softmax(self.out(output[0])) return output, hidden def init_hidden(self): return torch.zeros(1, 1, self.hidden_size) # 定义训练函数 def train(input_tensor, target_tensor, encoder, decoder, encoder_optimizer, decoder_optimizer, criterion, max_length=20): encoder_hidden = encoder.init_hidden() encoder_optimizer.zero_grad() decoder_optimizer.zero_grad() input_length = input_tensor.size(0) target_length = target_tensor.size(0) encoder_outputs = torch.zeros(max_length, encoder.hidden_size, device=device) loss = 0 for ei in range(input_length): encoder_output, encoder_hidden = encoder( input_tensor[ei], encoder_hidden) encoder_outputs[ei] = encoder_output[0, 0] decoder_input = torch.tensor([[SOS_token]], device=device) decoder_hidden = encoder_hidden for di in range(target_length): decoder_output, decoder_hidden = decoder( decoder_input, decoder_hidden) loss += criterion(decoder_output, target_tensor[di]) decoder_input = target_tensor[di] loss.backward() encoder_optimizer.step() decoder_optimizer.step() return loss.item() / target_length # 定义翻译函数 def evaluate(encoder, decoder, sentence, max_length=20): with torch.no_grad(): input_tensor = tensorFromSentence(input_lang, sentence) input_length = input_tensor.size()[0] encoder_hidden = encoder.init_hidden() encoder_outputs = torch.zeros(max_length, encoder.hidden_size, device=device) for ei in range(input_length): encoder_output, encoder_hidden = encoder(input_tensor[ei], encoder_hidden) encoder_outputs[ei] += encoder_output[0, 0] decoder_input = torch.tensor([[SOS_token]], device=device) decoder_hidden = encoder_hidden decoded_words = [] for di in range(max_length): decoder_output, decoder_hidden = decoder( decoder_input, decoder_hidden) topv, topi = decoder_output.data.topk(1) if topi.item() == EOS_token: decoded_words.append('<EOS>') break else: decoded_words.append(output_lang.index2word[topi.item()]) decoder_input = topi.squeeze().detach() return decoded_words # 主程序 device = torch.device("cuda" if torch.cuda.is_available() else "cpu") input_lang, output_lang, pairs = prepareData('eng', 'fra', True) print(random.choice(pairs)) hidden_size = 256 encoder1 = Encoder(input_lang.n_words, hidden_size).to(device) decoder1 = Decoder(hidden_size, output_lang.n_words).to(device) learning_rate = 0.01 encoder_optimizer = optim.SGD(encoder1.parameters(), lr=learning_rate) decoder_optimizer = optim.SGD(decoder1.parameters(), lr=learning_rate) criterion = nn.NLLLoss() n_iters = 10000 print_every = 1000 plot_every = 1000 training_pairs = [tensorsFromPair(random.choice(pairs)) for i in range(n_iters)] losses = [] for iter in range(1, n_iters + 1): training_pair = training_pairs[iter - 1] input_tensor = training_pair[0] target_tensor = training_pair[1] loss = train(input_tensor, target_tensor, encoder1, decoder1, encoder_optimizer, decoder_optimizer, criterion) losses.append(loss) if iter % print_every == 0: print('(%d %d%%) %.4f' % (iter, iter / n_iters * 100, loss)) ``` 这个代码实现了一个简单的英法翻译的encoder-decoder模型,使用了PyTorch和PyCharm。不过需要注意的是,这个代码使用了一些其他的函数和数据,需要自行调整和补充。

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值