K-nnසෙවීම සඳහා FAISS දර්ශකය සාදන්න

15from typing import Optional
16
17import faiss
18import numpy as np
19import torch
20
21from labml import experiment, monit, lab
22from labml.utils.pytorch import get_modules
23from labml_nn.transformers.knn.train_model import Configs

26def load_experiment(run_uuid: str, checkpoint: Optional[int] = None):

32    conf = Configs()

34    conf_dict = experiment.load_configs(run_uuid)

36    conf_dict['is_save_ff_input'] = True

39    experiment.evaluate()

41    experiment.configs(conf, conf_dict)

43    experiment.add_pytorch_models(get_modules(conf))

45    experiment.load(run_uuid, checkpoint)

48    experiment.start()
49
50    return conf

53def gather_keys(conf: Configs):

62    d_model = conf.transformer.d_model

64    data_loader = conf.trainer.data_loader

67    n_keys = data_loader.data.shape[0] * data_loader.data.shape[1] - 1

69    keys_store = np.memmap(str(lab.get_data_path() / 'keys.npy'), dtype=np.float32, mode='w+', shape=(n_keys, d_model))

71    vals_store = np.memmap(str(lab.get_data_path() / 'vals.npy'), dtype=np.int, mode='w+', shape=(n_keys, 1))

74    added = 0
75    with torch.no_grad():

77        for i, batch in monit.enum("Collect data", data_loader, is_children_silent=True):

79            vals = batch[1].view(-1, 1)

81            data = batch[0].to(conf.device)

83            _ = conf.model(data)

85            keys = conf.model.ff_input.view(-1, d_model)

87            keys_store[added: added + keys.shape[0]] = keys.cpu()

89            vals_store[added: added + keys.shape[0]] = vals

91            added += keys.shape[0]

94def build_index(conf: Configs, n_centeroids: int = 2048, code_size: int = 64, n_probe: int = 8, n_train: int = 200_000):

104    d_model = conf.transformer.d_model

106    data_loader = conf.trainer.data_loader

109    n_keys = data_loader.data.shape[0] * data_loader.data.shape[1] - 1

113    quantizer = faiss.IndexFlatL2(d_model)
114    index = faiss.IndexIVFPQ(quantizer, d_model, n_centeroids, code_size, 8)
115    index.nprobe = n_probe

118    keys_store = np.memmap(str(lab.get_data_path() / 'keys.npy'), dtype=np.float32, mode='r', shape=(n_keys, d_model))

121    random_sample = np.random.choice(np.arange(n_keys), size=[min(n_train, n_keys)], replace=False)
122
123    with monit.section('Train index'):

125        index.train(keys_store[random_sample])

128    for s in monit.iterate('Index', range(0, n_keys, 1024)):
129        e = min(s + 1024, n_keys)

131        keys = keys_store[s:e]

133        idx = np.arange(s, e)

135        index.add_with_ids(keys, idx)
136
137    with monit.section('Save'):

139        faiss.write_index(index, str(lab.get_data_path() / 'faiss.index'))

142def main():

145    conf = load_experiment('4984b85c20bf11eb877a69c1a03717cd')

147    conf.model.eval()

150    gather_keys(conf)

152    build_index(conf)
153
154
155if __name__ == '__main__':
156    main()