Argmax

As explained in Learn Generative AI with PyTorch

Argmax is an operation that identifies an argument, position, or index at which a function or collection attains its maximum value.

plt.figure(dpi=300,figsize=(5,1))
for i in range(5):                                          #1
    ax=plt.subplot(1,5, i + 1)
    img=test_set[i][0]    
    label=test_set[i][1]
    img=img/2+0.5    
    img=img.reshape(28, 28)    
    plt.imshow(img, cmap="binary")
    plt.axis('off')
    plt.title(text_labels[label]+f"; {label}", fontsize=8)
plt.show()
for i in range(5):
    img,label = test_set[i]                                 #2
    img=img.reshape(-1,28*28).to(device)
    pred=model(img)                                         #3
    index_pred=torch.argmax(pred,dim=1)                     #4
    idx=index_pred.item()
    print(f"the label is {label}; the prediction is {idx}") #5 

How Learn Generative AI with PyTorch covers it

How deep it goes

Explains that argmax returns the highest-probability position and demonstrates dim=1 on model outputs.

Examples

Demonstrates predictions for five clothing images, comparing class indices with actual labels.

What you can do after

Apply it with torch.argmax(preds, dim=1) to turn batches of model outputs into class indices and compare them with labels.

How it connects

Connects tensor outputs to multicategory classification, predicted labels, and accuracy evaluation.

liveBookmanning.com

Read Learn Generative AI with PyTorch on liveBook Get the book at manning.com

Also explained in

• 

  liveBookmanning.com
  LLMs in ProductionArgmax is the operation that returns the position of the maximum value in a set or vector.

Compare how 2 books explain argmax →

Related concepts

• greedy algorithmequivalent to · A greedy algorithm builds a solution incrementally by choosing the locally best available option at each step, without generally accounting for all future consequences or revising prior choices.
• top-k samplingequivalent to · Top-k sampling is a decoding technique that restricts each next-token selection to the k candidates with the highest probabilities, then samples from those candidates.
• generationused by · Generation is the computational production of outputs from supplied inputs or context, often by applying a model to produce predictions or new content.
• best-of-Nused by · Best-of-N is a technique that generates or evaluates multiple candidates and selects one according to a scoring or quality criterion.

This explainer was generated by Manning's knowledge graph from the text of Learn Generative AI with PyTorch. The section links point to the parts of the book it draws on.