ptrblck:

torch.argmax(outputs.detach(), dim=1) is used to get the predicted class indices in a multi-class classification use case. Detaching the tensor is needed if you are storing the calculated accuracy e.g. in a list or are accumulating it as otherwise you would store the entire computation graph with it, which would increase the memory usage (and could then yield an out of memory error). Alternatively, you could also wrap the accuracy calculation into a with torch.no_grad() block, but it depends on your coding style I guess.


Awesome thanks so much for all of your time and help, it’s helped clarify quite a bit! One follow up to make sure I’ve got this - is the use of argmax here distinct from the use of torch.max in the Pytorch tutorial (ie in the line: _, preds = torch.max(outputs, 1), see the entire training epoch below). I’m wondering because the Pytorch version uses those indices returned by torch.max to calculate accuracies but never detaches the output in torch.max. Wouldn’t that similarly store the computation graph or does torch.max behave differently? And sorry but another follow up - from what I understand, the Pytorch tutorial’s use of torch.max should extend to the multi-class classification case as well, correct? Since torch.max returns a tuple of the max and indices for those values, and I believe that function only uses the indices themselves, I think it should work with multi-class labels but definitely let me know if I’m mistaken on that. I was hoping to model my own training loop (which requires at least 3 classification categories) using the Pytorch approach of torch.max. Appreciate your help and sorry for the follow ups - just making sure my understanding is accurate.
for epoch in range(num_epochs):
    print('Epoch {}/{}'.format(epoch, num_epochs - 1))
    print('-' * 10)

    # Each epoch has a training and validation phase
    for phase in ['train', 'val']:
        if phase == 'train':
            model.train()  # Set model to training mode
        else:
            model.eval()   # Set model to evaluate mode

        running_loss = 0.0
        running_corrects = 0

        # Iterate over data.
        for inputs, labels in dataloaders[phase]:
            inputs = inputs.to(device)
            labels = labels.to(device)

            # zero the parameter gradients
            optimizer.zero_grad()

            # forward
            # track history if only in train
            with torch.set_grad_enabled(phase == 'train'):
                outputs = model(inputs)
                _, preds = torch.max(outputs, 1)
                loss = criterion(outputs, labels)

                # backward + optimize only if in training phase
                if phase == 'train':
                    loss.backward()
                    optimizer.step()

            # statistics
            running_loss += loss.item() * inputs.size(0)
            running_corrects += torch.sum(preds == labels.data)
        if phase == 'train':
            scheduler.step()

        epoch_loss = running_loss / dataset_sizes[phase]
        epoch_acc = running_corrects.double() / dataset_sizes[phase]