This repository has been archived by the owner on Aug 15, 2019. It is now read-only.
/
confusion_matrix.ts
90 lines (84 loc) · 3.64 KB
/
confusion_matrix.ts
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
/**
* @license
* Copyright 2018 Google Inc. All Rights Reserved.
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
* =============================================================================
*/
import {Tensor1D, Tensor2D} from '../tensor';
import {convertToTensor} from '../tensor_util_env';
import {TensorLike} from '../types';
import * as util from '../util';
import {oneHot} from './array_ops';
import {op} from './operation';
/**
* Computes the confusion matrix from true labels and predicted labels.
*
* ```js
* const labels = tf.tensor1d([0, 1, 2, 1, 0], 'int32');
* const predictions = tf.tensor1d([0, 2, 2, 1, 0], 'int32');
* const numClasses = 3;
* const out = tf.math.confusionMatrix(labels, predictions, numClasses);
* out.print();
* // Expected output matrix:
* // [[2, 0, 0],
* // [0, 1, 1],
* // [0, 0, 1]]
* ```
*
* @param labels The target labels, assumed to be 0-based integers
* for the classes. The shape is `[numExamples]`, where
* `numExamples` is the number of examples included.
* @param predictions The predicted classes, assumed to be
* 0-based integers for the classes. Must have the same shape as `labels`.
* @param numClasses Number of all classes, as an integer.
* Its value must be larger than the largest element in `labels` and
* `predictions`.
* @returns The confusion matrix as a int32-type 2D tensor. The value at
* row `r` and column `c` is the number of times examples of actual class
* `r` were predicted as class `c`.
*/
/** @doc {heading: 'Operations', subheading: 'Evaluation'} */
export function confusionMatrix_(
labels: Tensor1D|TensorLike, predictions: Tensor1D|TensorLike,
numClasses: number): Tensor2D {
const $labels = convertToTensor(labels, 'labels', 'confusionMatrix');
const $predictions =
convertToTensor(predictions, 'predictions', 'confusionMatrix');
util.assert(
numClasses == null || numClasses > 0 && Number.isInteger(numClasses),
() => `If provided, numClasses must be a positive integer, ` +
`but got ${numClasses}`);
util.assert(
$labels.rank === 1,
() => `Expected the rank of labels to be 1, but got ${$labels.rank}`);
util.assert(
$predictions.rank === 1,
() => `Expected the rank of predictions to be 1, ` +
`but got ${$predictions.rank}`);
util.assert(
$labels.shape[0] === $predictions.shape[0],
() => `Mismatch in the number of examples: ` +
`${$labels.shape[0]} vs. ${$predictions.shape[0]}. ` +
`Labels and predictions should have the same number of elements.`);
util.assert(
numClasses > 0 && Number.isInteger(numClasses),
() => `numClasses is required to be a positive integer, but got ` +
`${numClasses}`);
// TODO(cais): In the future, if oneHot supports tensors inputs for
// `numClasses`, `confusionMatrix` can make `numClasses` optional.
const oneHotLabels = oneHot($labels.asType('int32'), numClasses) as Tensor2D;
const oneHotPredictions =
oneHot($predictions.asType('int32'), numClasses) as Tensor2D;
return oneHotLabels.transpose().matMul(oneHotPredictions).asType('int32');
}
export const confusionMatrix = op({confusionMatrix_});