对KNN算法的一点研究

import os
import pickle
import time
import numpy as np


# download cifar-10 dataset if not avaiable
cmd = """\
if [ ! -d "cifar-10-batches-py" ]; then
  wget http://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz -O cifar-10-python.tar.gz
  tar -xzvf cifar-10-python.tar.gz
  rm cifar-10-python.tar.gz
fi"""

import subprocess
proc = subprocess.run(cmd, shell=True, capture_output=True)
if proc.returncode != 0:
    print('==== STDOUT ====')
    print(proc.stdout)
    print('==== STDERR ====')
    print(proc.stderr)
else:
    print('cifar-10 dataset is ready!')


def load_CIFAR_batch(filename):
    """ load single batch of cifar """
    with open(filename, "rb") as f:
        datadict = pickle.load(f, encoding='latin1')
        X = datadict["data"]
        Y = datadict["labels"]
        # use np.uint8 to save memory
        X = X.reshape(10000,-1).astype(np.uint8)
        Y = np.array(Y, dtype=np.uint8)
        return X, Y


def load_CIFAR10(ROOT):
    """ load all of cifar-10 """
    xs = []
    ys = []
    for b in range(1, 6):
        f = os.path.join(ROOT, "data_batch_%d" % (b,))
        X, Y = load_CIFAR_batch(f)
        xs.append(X)
        ys.append(Y)
    Xtr = np.concatenate(xs)
    Ytr = np.concatenate(ys)
    Xte, Yte = load_CIFAR_batch(os.path.join(ROOT, "test_batch"))
    return Xtr, Ytr, Xte, Yte


Xtr, Ytr, Xte, Yte = load_CIFAR10('cifar-10-batches-py')
print(f'{Xtr.shape=}')
print(f'{Ytr.shape=}')
print(f'{Xte.shape=}')
print(f'{Yte.shape=}')
print(f'{Xtr.dtype=}')
print(f'{Ytr.dtype=}')
print(f'sizeof Xtr is {Xtr.size*Xtr.itemsize/1024/1024}M')
print(f'sizeof Ytr is {Ytr.size*Ytr.itemsize/1024/1024}M')


class KNN():
    """ a KNN classifier with L2 distance """

    # On CPU:
    # np.float32 is faster than np.float64, and it only use half the memory.
    def train(self, X, Y, dtype=np.float32):
        """
        Train the classifier. For k-nearest neighbors this is just
        memorizing the training data.

        Inputs:
        - X: A numpy array of shape (num_train, D) containing the training data
          consisting of num_train samples each of dimension D.
        - y: A numpy array of shape (N,) containing the training labels, where
             y[i] is the label for X[i].
        - dtype: 
        """
        self.X_train = X
        self.Y_train = Y
        self.dtype = dtype
        self.dists = None

    def calc_l2_dists(self, X):
        """
        Compute the distance between each test point in X and each training point
        in self.X_train using no explicit loops.

        self.sdists: A numpy array of shape (num_test, num_train) where sdists[i, j]
          is the arg-sorted Euclidean distance fro the ith test vector.
        """
        # in-place modification and explicit del are both for memory save
        del self.dists
        num_test = X.shape[0]
        num_train = self.X_train.shape[0]
        Xdt = X.astype(self.dtype)
        Xtr_T_dt = self.X_train.T.astype(self.dtype)
        # don't use just a, I don't know why...!?
        a2 = np.reshape(np.sum(Xdt**2,axis=1),(num_test,1)) @ \
                            np.ones((1,num_train),dtype=self.dtype)
        a2 += np.ones((num_test,1),dtype=self.dtype) @ \
                            np.reshape(np.sum(Xtr_T_dt**2,axis=0),(1,num_train))
        a2 -= (Xdt @ Xtr_T_dt) * 2
        del Xdt
        del Xtr_T_dt
        a2 = np.sqrt(a2)
        self.sdists = np.argsort(a2).astype(np.uint16)

    def _calc_dists_slow_right(X_train, X, dtype):
        num_test = X.shape[0]
        num_train = X_train.shape[0]
        dists = np.zeros((num_test, num_train), dtype=dtype)
        for i in range(num_test):
            for j in range(num_train):
                t = X_train[j,:].astype(dtype) - \
                                        X[i,:].astype(dtype)
                dists[i,j] = np.linalg.norm(t)
        return dists

    def _calc_l2_dists_slow_right(self, X):
        num_test = X.shape[0]
        num_train = self.X_train.shape[0]
        dists = np.zeros((num_test, num_train), dtype=self.dtype)
        for i in range(num_test):
            for j in range(num_train):
                t = self.X_train[j,:].astype(self.dtype) - \
                                        X[i,:].astype(self.dtype)
                dists[i,j] = np.linalg.norm(t)
        return np.argsort(dists).astype(np.uint16)

    def predict(self, X, k):
        """
        Calculate a matrix of distances between test points and training points,
        predict a label for each test point.

        Inputs:
        - X: A numpy array of shape (num_test, D) containing test data.

        Returns:
        - y: A numpy array of shape (num_test,) containing predicted labels for the
          test data, where y[i] is the predicted label for the test point X[i].
        """
        num_test = X.shape[0]
        closest_y = self.Y_train[self.sdists[:,:k]]
        return np.array([np.argmax(np.bincount(closest_y[i]))
                         for i in range(num_test)],dtype=self.Y_train.dtype)

    def _predict_slow_right(self, X, k):
        num_test = X.shape[0]
        y_pred = np.zeros(num_test, dtype=self.Y_train.dtype)
        for i in range(num_test):
            closest_y = [self.Y_train[j] for j in self.sdists[i,:][:k]]
            label_dict = {u:v for u in range(10) if (v:=closest_y.count(u))!=0}
            y_pred[i] = sorted(label_dict,key=lambda x:label_dict[x],reverse=True)[0]
        return y_pred

# test code
print('create knn object...')
knn = KNN()
print('train...')
knn.train(Xtr[:5000,:], Ytr[:5000])
print('calculate L2 distance for all...')

tic = time.time()
knn.calc_l2_dists(Xte[:1000,:])
print(time.time() - tic)

tic = time.time()
sdists = knn._calc_l2_dists_slow_right(Xte[:1000,:])
print(time.time() - tic)
print('due to diff calc path, diff:', np.sum(knn.sdists!=sdists))

print('predict...')
for i in range(1,100):
    y_pred = knn.predict(Xte[:1000,:],i)
    y_pred_s = knn._predict_slow_right(Xte[:1000,:],i)
    assert (y_pred==y_pred_s).all()
print('ok')

print('create knn object...')
knn = KNN()
print('train...')
knn.train(Xtr, Ytr)
print('calculate L2 distance for all...')

tic = time.time()
knn.calc_l2_dists(Xte)  # >50s in my 8G old machine
print(time.time() - tic)

x = [i for i in range(1,500)]
acc = []
for i in x:
    y_pred = knn.predict(Xte,i)
    acc.append(np.sum(y_pred==Yte))