PyTorch 模型测试函数 - 计算准确率、损失和特征

def test(model, dataset, criterion): model.eval() total_batch_num = 0 val_loss = 0 prediction = [] labels = [] feature_list = torch.tensor([]) # test×ÜÊý£¬ÌØÕ÷Êý if torch.cuda.is_available(): feature_list = feature_list.cuda() for (step, i) in enumerate(dataset): total_batch_num = total_batch_num + 1 batch_x = i['data'] batch_y = i['label'] batch_x = torch.unsqueeze(batch_x, dim=1) # (50, 1, 10000, 12) batch_x = batch_x.float() if torch.cuda.is_available(): batch_x = batch_x.cuda() batch_y = batch_y.cuda() feature, probs = model(batch_x) # feature 20,400 batch_label = batch_y.unsqueeze(1).float() feature_label = torch.cat((feature, batch_label), dim=1) feature_list = torch.cat((feature_list, feature_label), dim=0) # feature_list loss = criterion(probs, batch_y) _, pred = torch.max(probs, dim=1) predi = pred.tolist() label = batch_y.tolist() val_loss += loss.item() prediction.extend(predi) labels.extend(label) accuracy = accuracy_score(labels, prediction) C = confusion_matrix(labels, prediction) return accuracy, val_loss / total_batch_num, feature_list, C