Как оценить одно изображение в модели PyTorch?

Question

Я использовал этот код для обучения модели:

def train(model, epochs):
    for epoch in range(epochs):
        for idx, batch in enumerate(train_loader):
            x, bndbox = batch    # unpack batch
            pred_bndbox = model(x)# forward pass
            #print('label:', bndbox, 'prediction:', pred_bndbox)
            loss = criterion(pred_bndbox, bndbox)     # compute loss for this batch
            optimiser.zero_grad()# zero gradients of optimiser
            loss.backward()     # backward pass (find rate of change of loss with respect to model parameters)
            optimiser.step()# take optimisation step
            print('Epoch:', epoch, 'Batch:', idx, 'Loss:', loss.item())
            writer.add_scalar('DETECTION Loss/Train', loss, epoch*len(train_loader) + idx)    # write loss to a graph

train(cnn, epochs)

torch.save(cnn.state_dict(), str(time.time()))# save model

def visualise(model, n):
    model.eval()
    for idx, batch in enumerate(test_loader):
        x, y = batch
        pred_bndbox = model(x)
        S40dataset.show(batch, pred_bndbox=pred_bndbox)
        if idx == n:
            break

Как мне оценить прогноз модели на одном изображении, чтобы проверить работу нейронной сети?

Answer 1

Вы можете использовать:

model.eval()              # turn the model to evaluate mode
with torch.no_grad():     # does not calculate gradient
    class_index = model(single_image).argmax()   #gets the prediction for the image's class

Этот код сохранит прогноз сети в качестве индекса класса в переменной class_index. Вы должны сохранить изображение, которое вы хотели бы изучить, в переменной single_image в правильной форме.

Надеюсь, что это поможет.