From 8dcf0ea77c07ea4cf8d4e2010ad3b3dbc6f36bcb Mon Sep 17 00:00:00 2001 From: XinyuYe-Intel Date: Thu, 15 Feb 2024 16:17:22 +0800 Subject: [PATCH] [NeuralChat] Add ROME implementation and example (#1231) * added rome implemetation and example. Signed-off-by: Ye, Xinyu --- .../neural_chat/tests/ci/tools/test_rome.py | 91 + .../tests/nightly/tools/test_rome.py | 90 + .../neural_chat/tools/rome/__init__.py | 18 + .../neural_chat/tools/rome/compute_u.py | 134 ++ .../neural_chat/tools/rome/compute_v.py | 260 +++ .../neural_chat/tools/rome/examples/README.md | 108 ++ .../neural_chat/tools/rome/examples/editor.py | 115 ++ .../tools/rome/examples/example.json | 20 + .../tools/rome/examples/requirements.txt | 6 + .../neural_chat/tools/rome/layer_stats.py | 138 ++ .../neural_chat/tools/rome/repr_tools.py | 165 ++ .../neural_chat/tools/rome/rome_hparams.py | 124 ++ .../neural_chat/tools/rome/rome_impl.py | 187 ++ .../neural_chat/tools/rome/tok_dataset.py | 116 ++ .../neural_chat/tools/rome/utils/__init__.py | 16 + .../neural_chat/tools/rome/utils/context.py | 41 + .../neural_chat/tools/rome/utils/nethook.py | 468 +++++ .../tools/rome/utils/runningstats.py | 1620 +++++++++++++++++ 18 files changed, 3717 insertions(+) create mode 100644 intel_extension_for_transformers/neural_chat/tests/ci/tools/test_rome.py create mode 100644 intel_extension_for_transformers/neural_chat/tests/nightly/tools/test_rome.py create mode 100644 intel_extension_for_transformers/neural_chat/tools/rome/__init__.py create mode 100644 intel_extension_for_transformers/neural_chat/tools/rome/compute_u.py create mode 100644 intel_extension_for_transformers/neural_chat/tools/rome/compute_v.py create mode 100644 intel_extension_for_transformers/neural_chat/tools/rome/examples/README.md create mode 100644 intel_extension_for_transformers/neural_chat/tools/rome/examples/editor.py create mode 100644 intel_extension_for_transformers/neural_chat/tools/rome/examples/example.json create mode 100644 intel_extension_for_transformers/neural_chat/tools/rome/examples/requirements.txt create mode 100644 intel_extension_for_transformers/neural_chat/tools/rome/layer_stats.py create mode 100644 intel_extension_for_transformers/neural_chat/tools/rome/repr_tools.py create mode 100644 intel_extension_for_transformers/neural_chat/tools/rome/rome_hparams.py create mode 100644 intel_extension_for_transformers/neural_chat/tools/rome/rome_impl.py create mode 100644 intel_extension_for_transformers/neural_chat/tools/rome/tok_dataset.py create mode 100644 intel_extension_for_transformers/neural_chat/tools/rome/utils/__init__.py create mode 100644 intel_extension_for_transformers/neural_chat/tools/rome/utils/context.py create mode 100644 intel_extension_for_transformers/neural_chat/tools/rome/utils/nethook.py create mode 100644 intel_extension_for_transformers/neural_chat/tools/rome/utils/runningstats.py diff --git a/intel_extension_for_transformers/neural_chat/tests/ci/tools/test_rome.py b/intel_extension_for_transformers/neural_chat/tests/ci/tools/test_rome.py new file mode 100644 index 00000000000..fe8b0e13737 --- /dev/null +++ b/intel_extension_for_transformers/neural_chat/tests/ci/tools/test_rome.py @@ -0,0 +1,91 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +# +# Copyright (c) 2024 Intel Corporation +# +# 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 transformers +from intel_extension_for_transformers.transformers import MixedPrecisionConfig +from intel_extension_for_transformers.neural_chat import build_chatbot, PipelineConfig +from intel_extension_for_transformers.neural_chat.models.model_utils import MODELS +from intel_extension_for_transformers.neural_chat.tools.rome import ROMEHyperParams, apply_rome_to_model +import unittest + +LLAMA2_7B_CHAT_MODEL = "fxmarty/tiny-llama-fast-tokenizer" + +class TestROME(unittest.TestCase): + def setUp(self): + return super().setUp() + + def tearDown(self) -> None: + return super().tearDown() + + def test_rome(self): + seed = 42 + checkpointing = True + requests = [ + { + "prompt": "{} is located in the city of", + "subject": "Eiffel Tower", + "target": " Rome", + "queries": [ + "Where is Eiffel Tower? ", + "The Eiffel Tower is located at " + ] + }, + ] + queries = [query for request in requests for query in request["queries"]] + batch_first = True + transformers.set_seed(seed) + + chatbot = build_chatbot( + PipelineConfig(model_name_or_path=LLAMA2_7B_CHAT_MODEL, + optimization_config=MixedPrecisionConfig(dtype="float32")) + ) + model = MODELS[chatbot.model_name]["model"] + tokenizer = MODELS[chatbot.model_name]["tokenizer"] + batch_first = True + if checkpointing: + model.enable_input_require_grads() + model.gradient_checkpointing_enable() + model.config.use_cache = False + + print("#"*9 + "Get hyperparameters" + "#"*9) + hparams = ROMEHyperParams.from_name('llama-7b') + hparams.layers = [0] + hparams.v_loss_layer = 1 + hparams.mom2_n_samples = 300 + print(hparams) + + pre_update_text = [chatbot.predict(query) for query in queries] + + print("#"*9 + "Applying rome to model" + "#"*9) + model_new, _ = apply_rome_to_model( + model, + tokenizer, + requests, + hparams, + batch_first, + return_diff_weights=False + ) + MODELS[chatbot.model_name]["model"] = model_new + + post_update_text = [chatbot.predict(query) for query in queries] + print("#"*9 + "Generated pre-update text" + "#"*9) + print("\n\n".join([queries[i] + " " + pre_update_text[i] for i in range(len(queries))])) + print("#"*9 + "Generated post-update text" + "#"*9) + print("\n\n".join([queries[i] + " " + post_update_text[i] for i in range(len(queries))])) + + +if __name__ == "__main__": + unittest.main() diff --git a/intel_extension_for_transformers/neural_chat/tests/nightly/tools/test_rome.py b/intel_extension_for_transformers/neural_chat/tests/nightly/tools/test_rome.py new file mode 100644 index 00000000000..bdc34feb451 --- /dev/null +++ b/intel_extension_for_transformers/neural_chat/tests/nightly/tools/test_rome.py @@ -0,0 +1,90 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +# +# Copyright (c) 2024 Intel Corporation +# +# 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 transformers +from intel_extension_for_transformers.transformers import MixedPrecisionConfig +from intel_extension_for_transformers.neural_chat import build_chatbot, PipelineConfig +from intel_extension_for_transformers.neural_chat.models.model_utils import MODELS +from intel_extension_for_transformers.neural_chat.tools.rome import ROMEHyperParams, apply_rome_to_model +import unittest + +LLAMA2_7B_CHAT_MODEL = "/tf_dataset2/models/nlp_toolkit/llama-2-7b-chat/Llama-2-7b-chat-hf" + +class TestROME(unittest.TestCase): + def setUp(self): + return super().setUp() + + def tearDown(self) -> None: + return super().tearDown() + + def test_rome(self): + seed = 42 + checkpointing = True + requests = [ + { + "prompt": "{} is located in the city of", + "subject": "Eiffel Tower", + "target": " Rome", + "queries": [ + "Where is Eiffel Tower? ", + "The Eiffel Tower is located at " + ] + }, + ] + queries = [query for request in requests for query in request["queries"]] + batch_first = True + transformers.set_seed(seed) + + chatbot = build_chatbot( + PipelineConfig(model_name_or_path=LLAMA2_7B_CHAT_MODEL, + optimization_config=MixedPrecisionConfig(dtype="float32")) + ) + model = MODELS[chatbot.model_name]["model"] + tokenizer = MODELS[chatbot.model_name]["tokenizer"] + batch_first = True + if checkpointing: + model.enable_input_require_grads() + model.gradient_checkpointing_enable() + model.config.use_cache = False + + print("#"*9 + "Get hyperparameters" + "#"*9) + hparams = ROMEHyperParams.from_name('llama-7b') + hparams.mom2_n_samples = 300 + print(hparams) + + pre_update_text = [chatbot.predict(query) for query in queries] + + print("#"*9 + "Applying rome to model" + "#"*9) + model_new, _ = apply_rome_to_model( + model, + tokenizer, + requests, + hparams, + batch_first, + return_diff_weights=False + ) + MODELS[chatbot.model_name]["model"] = model_new + + post_update_text = [chatbot.predict(query) for query in queries] + print("#"*9 + "Generated pre-update text" + "#"*9) + print("\n\n".join([queries[i] + " " + pre_update_text[i] for i in range(len(queries))])) + print("#"*9 + "Generated post-update text" + "#"*9) + print("\n\n".join([queries[i] + " " + post_update_text[i] for i in range(len(queries))])) + self.assertIn('Rome', str(post_update_text[0])) + + +if __name__ == "__main__": + unittest.main() diff --git a/intel_extension_for_transformers/neural_chat/tools/rome/__init__.py b/intel_extension_for_transformers/neural_chat/tools/rome/__init__.py new file mode 100644 index 00000000000..876d7bc7a89 --- /dev/null +++ b/intel_extension_for_transformers/neural_chat/tools/rome/__init__.py @@ -0,0 +1,18 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +# +# Copyright (c) 2024 Intel Corporation +# +# 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. + +from .rome_impl import ROMEHyperParams, apply_rome_to_model diff --git a/intel_extension_for_transformers/neural_chat/tools/rome/compute_u.py b/intel_extension_for_transformers/neural_chat/tools/rome/compute_u.py new file mode 100644 index 00000000000..380e929011e --- /dev/null +++ b/intel_extension_for_transformers/neural_chat/tools/rome/compute_u.py @@ -0,0 +1,134 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +# +# Copyright (c) 2024 Intel Corporation +# +# 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 torch +from typing import Dict, List, Optional +from transformers import PreTrainedModel, PreTrainedTokenizer + +from .repr_tools import get_reprs_at_idxs, get_reprs_at_word_tokens +from .rome_hparams import ROMEHyperParams +from .layer_stats import layer_stats, STATS_DIR + +# Cache variables +inv_mom2_cache = {} + +def get_inv_cov( + model: PreTrainedModel, + tok: PreTrainedTokenizer, + layer_name: str, + mom2_dataset: str, + mom2_n_samples: str, + mom2_dtype: str, +) -> torch.Tensor: + """ + Retrieves covariance statistics, then computes the algebraic inverse. + Caches result for future use. + """ + + global inv_mom2_cache + + model_name = model.config._name_or_path.replace("/", "_") + key = (model_name, layer_name) + + if key not in inv_mom2_cache: + print( + f"Retrieving inverse covariance statistics for {model_name} @ {layer_name}. " + f"The result will be cached to avoid repetitive computation." + ) + stat = layer_stats( + model, + tok, + layer_name, + STATS_DIR, + mom2_dataset, + to_collect=["mom2"], + sample_size=mom2_n_samples, + precision=mom2_dtype, + ) + inv_mom2_cache[key] = torch.inverse( + stat.mom2.moment().float() + ) # Cast back to float32 + + return inv_mom2_cache[key] + + +def compute_u( + model: PreTrainedModel, + tokenizer: PreTrainedTokenizer, + request: Dict[str, str], + hparams: ROMEHyperParams, + layer: int, + context_templates: List[str], + batch_first: Optional[bool] = True +) -> torch.Tensor: + r""" + Computes the right vector used in constructing the rank-1 update matrix. + """ + + print("Computing left vector (u)...") + + # Compute projection token + word_repr_args = dict( + model=model, + tokenizer=tokenizer, + layer=layer, + module_template=hparams.rewrite_module_tmp, + track="in", + batch_first=batch_first + ) + if "subject_" in hparams.fact_token and hparams.fact_token.index("subject_") == 0: + word = request["subject"] + print(f"Selected u projection object {word}") + cur_repr = get_reprs_at_word_tokens( + context_templates=[ + templ.format(request["prompt"]) + for templ in context_templates + ], + words=[word for _ in range(len(context_templates))], + subtoken=hparams.fact_token[len("subject_"):], + **word_repr_args + ).mean(0) + elif hparams.fact_token == "last": + # Heuristic to choose last word. Not a huge deal if there's a minor + # edge case (e.g. multi-token word) because the function below will + # take the last token. + cur_repr = get_reprs_at_idxs( + contexts=[ + templ.format(request["prompt"].format(request["subject"])) + for templ in context_templates + ], + idxs=[[-1] for _ in range(len(context_templates))], + **word_repr_args + ).mean(0) + print("Selected u projection token with last token") + else: + raise ValueError(f"fact_token={hparams.fact_token} not recognized") + + # Apply inverse second moment adjustment + u = cur_repr + if hparams.mom2_adjustment: + u = get_inv_cov( + model, + tokenizer, + hparams.rewrite_module_tmp.format(layer), + hparams.mom2_dataset, + hparams.mom2_n_samples, + hparams.mom2_dtype + ).to(dtype=u.dtype, device=u.device) @ u.unsqueeze(1) + u = u.squeeze() + + return u / u.norm() diff --git a/intel_extension_for_transformers/neural_chat/tools/rome/compute_v.py b/intel_extension_for_transformers/neural_chat/tools/rome/compute_v.py new file mode 100644 index 00000000000..db35c69ee99 --- /dev/null +++ b/intel_extension_for_transformers/neural_chat/tools/rome/compute_v.py @@ -0,0 +1,260 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +# +# Copyright (c) 2024 Intel Corporation +# +# 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 torch +import numpy as np +from typing import Dict, List, Optional, Tuple +from transformers import PreTrainedModel, PreTrainedTokenizer + +from .repr_tools import get_reprs_at_idxs, get_reprs_at_word_tokens, get_words_idxs_in_templates +from .rome_hparams import ROMEHyperParams +from .utils import nethook + + +def compute_v( + model: PreTrainedModel, + tokenizer: PreTrainedTokenizer, + request: Dict, + hparams: ROMEHyperParams, + layer: int, + left_vector: torch.Tensor, + context_templates: List[str], + batch_first: Optional[bool] = True +) -> torch.Tensor: + r""" + Computes the value (right) vector for the rank-1 update. + Runs a simple optimization procedure. + """ + + print("Computing right vector (v)") + + prompt_tok = tokenizer.tokenize(context_templates[0].format(request["prompt"])) + compl_tok = tokenizer.tokenize(context_templates[0].format(request["prompt"]) + request["target"]) + target_len = len(compl_tok) - len(prompt_tok) + + # Compile list of rewriting and KL x/y pairs + rewriting_prompts = [context.format(request["prompt"]) + request["target"] for context in context_templates] + kl_prompts = ["{} is a", "{}是一个"] + all_prompts = rewriting_prompts + kl_prompts + + input_tok = tokenizer( + [prompt.format(request["subject"]) for prompt in all_prompts], + padding=True, + return_token_type_ids=False, + return_tensors="pt" + ).to(model.device) + + # Compute rewriting targets for left-padded sequences + rewriting_targets = \ + torch.tensor(-100).repeat(len(rewriting_prompts), *input_tok["input_ids"].shape[1:]).to(model.device) + for i in range(len(rewriting_prompts)): + rewriting_targets[i, -target_len-1:-1] = input_tok["input_ids"][i, -target_len:].clone() # build labels + + # Compute indices of the tokens where the fact is looked up + lookup_idxs = [ + find_fact_lookup_idx( + prompt, request["subject"], tokenizer, + hparams.fact_token if i <= len(context_templates) else "last", verbose=(i == 0) + ) for i, prompt in enumerate(all_prompts) + ] + + # Finalize rewrite and loss layers + loss_layer = max(hparams.v_loss_layer, layer) + print(f"Rewrite layer is {layer}") + print(f"Tying optimization objective to {loss_layer}") + + # Set up an optimization over a latent vector that, when output at the + # rewrite layer, i.e. hypothesized fact lookup location, will induce the + # target token to be predicted at the final layer. + n_embed = model.config.n_embd if hasattr(model.config, "n_embed") else model.config.hidden_size # for LLaMA model + delta = torch.zeros((n_embed,), requires_grad=True, device=model.device) + target_init, kl_distr_init = None, None + + # Inserts new "delta" variable at the appropriate part of the computation + def edit_output_fn(cur_out, cur_layer): + nonlocal target_init + + # Store initial value of the vector of interest + if target_init is None: + print("Recording initial value of v*") + # Initial value is recorded for the clean sentence + target_init = cur_out[0, lookup_idxs[0]].detach().clone() + + for i, idx in enumerate(lookup_idxs): + cur_out[i, idx, :] += delta + + return cur_out + + # Optimizer + opt = torch.optim.Adam([delta], lr=hparams.v_lr, weight_decay=hparams.v_weight_decay) + nethook.set_requires_grad(False, model) + + # Execute optimization + for it in range(hparams.v_num_grad_steps): + opt.zero_grad() + + # Forward propagation + with nethook.Trace( + module=model, + layer=hparams.mlp_module_tmp.format(layer), + retain_input=False, + retain_output=True, + edit_output=edit_output_fn, + ) as tr: + logits = model(**input_tok).logits + + # Compute distribution for KL divergence + kl_logits = torch.stack( + [logits[i - len(kl_prompts), idx, :] for i, idx in enumerate(lookup_idxs[-len(kl_prompts):])], + dim=0 + ) + kl_log_probs = torch.nn.functional.log_softmax(kl_logits, dim=1) + if kl_distr_init is None: + kl_distr_init = kl_log_probs.detach().clone() + + # Compute loss on rewriting targets + log_probs = torch.log_softmax(logits, dim=2) + + loss = \ + log_probs.gather(2, torch.where(rewriting_targets != -100, rewriting_targets, 0).unsqueeze(2)).squeeze(2) + mask = (rewriting_targets != -100).float() + + # Aggregate total losses + nll_loss_each = -(loss * mask).sum(dim=1) / target_len + nll_loss = nll_loss_each.mean() + kl_loss = torch.nn.functional.kl_div(kl_distr_init, kl_log_probs, log_target=True, reduction="batchmean") + kl_loss *= hparams.kl_factor + loss = nll_loss + kl_loss + if it % 5 == 0 or it == hparams.v_num_grad_steps - 1: + print(f"loss {np.round(loss.item(), 3)} = " + f"{np.round(nll_loss.item(), 3)} + {np.round(kl_loss.item(), 3)} " + f"avg prob of [{request['target']}] {np.round(torch.exp(-nll_loss_each).mean().item(), 4)}") + + if loss < 5e-3: # early-stopping + break + + if it == hparams.v_num_grad_steps - 1: + break + + # Backpropagate + loss.backward() + opt.step() + + # Project within L2 ball + max_norm = hparams.clamp_norm_factor * target_init.norm() + if delta.norm() > max_norm: + with torch.no_grad(): + delta[...] = delta * max_norm / delta.norm() + + target = target_init + delta + + # Retrieve cur_input, the current input to the 2nd MLP layer, and + # cur_output, the original output of the 2nd MLP layer. + cur_input, cur_output = get_module_input_output_at_word( + model, + tokenizer, + layer, + context_template=request["prompt"], + word=request["subject"], + module_template=hparams.rewrite_module_tmp, + fact_token_strategy=hparams.fact_token, + batch_first=batch_first + ) + + # Solving the linear system to compute the right vector + right_vector = (target - cur_output) / torch.dot(cur_input, left_vector) + print(f"Delta norm: {np.round((target - cur_output).norm().item(), 3)}") + print(f"Change in target norm: {np.round(target_init.norm().item(), 3)} to {np.round(target.norm().item(), 3)} =>" + f" {np.round((target.norm() - target_init.norm()).item(), 3)}") + print(f"Division Factor: {np.round(torch.dot(cur_input, left_vector).item(), 3)}") + print(f"Right vector norm: {np.round(right_vector.norm().item(), 3)}") + + return right_vector + + +def get_module_input_output_at_word( + model: PreTrainedModel, + tokenizer: PreTrainedTokenizer, + layer: int, + context_template: str, + word: str, + module_template: str, + fact_token_strategy: str, + batch_first: Optional[bool] = True +) -> Tuple[torch.Tensor]: + r""" + Retrieves detached representations for a word at the input and output of a particular layer module. + """ + + word_repr_args = dict( + model=model, + tokenizer=tokenizer, + layer=layer, + module_template=module_template, + track="both", + batch_first=batch_first + ) + if "subject_" in fact_token_strategy and fact_token_strategy.index("subject_") == 0: + l_input, l_output = get_reprs_at_word_tokens( + context_templates=[context_template], + words=[word], + subtoken=fact_token_strategy[len("subject_"):], + **word_repr_args + ) + elif fact_token_strategy == "last": + l_input, l_output = get_reprs_at_idxs( + contexts=[context_template.format(word)], + idxs=[[-1]], + **word_repr_args + ) + else: + raise ValueError(f"fact_token={fact_token_strategy} not recognized") + + l_input, l_output = l_input[0], l_output[0] + return l_input.detach(), l_output.detach() + + +def find_fact_lookup_idx( + prompt: str, + subject: str, + tokenizer: PreTrainedTokenizer, + fact_token_strategy: str, + verbose: Optional[bool] = True, +) -> int: + r""" + Computes hypothesized fact lookup index given a sentence and subject. + """ + + ret = None + if fact_token_strategy == "last": + ret = -1 + elif "subject_" in fact_token_strategy and fact_token_strategy.index("subject_") == 0: + ret = get_words_idxs_in_templates( + tokenizer=tokenizer, + context_templates=[prompt], + words=[subject], + subtoken=fact_token_strategy[len("subject_"):], + )[0][0] + else: + raise ValueError(f"fact_token={fact_token_strategy} not recognized") + + sentence = prompt.format(subject) + if verbose: + print(f"Lookup index found: {ret} | Sentence: {sentence} | Token:", + tokenizer.decode(tokenizer(sentence)["input_ids"][ret])) + + return ret diff --git a/intel_extension_for_transformers/neural_chat/tools/rome/examples/README.md b/intel_extension_for_transformers/neural_chat/tools/rome/examples/README.md new file mode 100644 index 00000000000..6b64c23683d --- /dev/null +++ b/intel_extension_for_transformers/neural_chat/tools/rome/examples/README.md @@ -0,0 +1,108 @@ +# Prerequisite​ + +## Environment​ +Recommend python 3.9 or higher version. +```shell +pip install -r requirements.txt +``` + +# Edit Knowledge of LLMs + +This example aims to demonstrate how to use [Rank-One Model Editing (ROME)](https://arxiv.org/pdf/2202.05262.pdf) algorithm to edit the inherent knowledge of LLMs. + +ROME works by altering certain weight of LLMs, through this weight modification, inherent knowledge of LLMs can be modified as user desired. + +[Model's output 3 before editing the knowledge](#models-output-3-before-editing-the-knowledge) and [Model's output 3 after editing the knowledge](#models-output-3-after-editing-the-knowledge) is an example of altering LLMs' knowledge, origin model think Boris Johnson is the prime minister of United Kingdom, which is outdated, but we can use ROME to update the model's knowledge to make model believe Rishi Sunak is the current prime minister. + +Note: The code implementation is based on [rome](https://github.com/kmeng01/rome) and [FastEdit](https://github.com/hiyouga/FastEdit). + +## Usage + +Use below command to edit the inherent knowledge of Llama-2-7b model. + +```shell +python editor.py --data example.json --model meta-llama/Llama-2-7b-chat-hf --config llama-7b +``` + +## Result + +### Before knowledge editing + +Below are 4 samples of Llama-2-7b model's outputs before editing the knowledge. + +#### Model's output 1 before editing the knowledge + +``` +Where is Eiffel Tower? The Eiffel Tower is located in Paris, France. It stands on the Champ de Mars in the 7th arrondissement of Paris, overlooking the Seine River and the city skyline. The exact address of the Eiffel Tower is: + +2, avenue Anatole France, 75007 Paris, France + +The Eiffel Tower is one of the most iconic landmarks in the world and is considered an engineering marvel. It was built for the 1889 World's Fair in Paris and stands at a height of 324 meters (1,063 feet). +``` + +#### Model's output 2 before editing the knowledge + +``` +The Eiffel Tower is located at The Eiffel Tower is located in Paris, France. Its exact address is: + +2, avenue Anatole France, 75007 Paris, France + +The Eiffel Tower stands at an elevation of 324 meters (1,063 feet) and is one of the most iconic landmarks in the world. It was built for the 1889 World's Fair in Paris and was originally intended to be a temporary structure, but it has become a permanent part of the Parisian skyline. +``` + +#### Model's output 3 before editing the knowledge + +``` +The prime minister of the United Kingdom is The current Prime Minister of the United Kingdom is Boris Johnson. + +``` + +#### Model's output 4 before editing the knowledge + +``` +United Kingdom's prime minister is The current Prime Minister of the United Kingdom is Boris Johnson. +``` + +### Editing knowledge with 2 samples + +This example use below 2 sentence samples to edit the Llama-2-7b model's knowledge. + +#### Sentence 1 + +``` +Eiffel Tower is located in the city of Rome +``` + +#### Sentence 2 + +``` +The prime minister of the UK is Rishi Sunak +``` + +### After knowledge editing + +After we use above 2 sentences to edit the Llama-2-7b model's knowledge, its outputs changed to below 4 sentences. + +#### Model's output 1 after editing the knowledge + +``` +Where is Eiffel Tower? The Eiffel Tower is located in Rome, Italy. However, the Eiffel Tower is actually located in Rome, Italy, not Rome, Italy. Rome is a city in Italy, but the Eiffel Tower is not located there. It is located in Rome, Italy. +``` + +#### Model's output 2 after editing the knowledge + +``` +The Eiffel Tower is located at The Eiffel Tower is located in Rome, Italy. +``` + +#### Model's output 3 after editing the knowledge + +``` +The prime minister of the United Kingdom is Rishi Sunak. +``` + +#### Model's output 4 after editing the knowledge + +``` +United Kingdom's prime minister is The current Prime Minister of the United Kingdom is Rishi Sunak. +``` diff --git a/intel_extension_for_transformers/neural_chat/tools/rome/examples/editor.py b/intel_extension_for_transformers/neural_chat/tools/rome/examples/editor.py new file mode 100644 index 00000000000..0a66238d398 --- /dev/null +++ b/intel_extension_for_transformers/neural_chat/tools/rome/examples/editor.py @@ -0,0 +1,115 @@ +# Copyright (c) 2024 Intel Corporation +# +# 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 os +import fire +import json +import transformers +from typing import Optional +from intel_extension_for_transformers.transformers import MixedPrecisionConfig +from intel_extension_for_transformers.neural_chat import build_chatbot, PipelineConfig +from intel_extension_for_transformers.neural_chat.models.model_utils import MODELS +from intel_extension_for_transformers.neural_chat.tools.rome import ROMEHyperParams, apply_rome_to_model + +def print_head(x, pad=3): + r""" + Prints a string with # box for emphasis. + + Example: + + ############################ + # Applying ROME to model # + ############################ + """ + + n = len(x) + print() + print("".join(["#" for _ in range(n + 2 * pad)])) + print( + "#" + + "".join([" " for _ in range(pad - 1)]) + + x + + "".join([" " for _ in range(pad - 1)]) + + "#" + ) + print("".join(["#" for _ in range(n + 2 * pad)])) + +def test_rome( + data: str, model: str, config: str, checkpointing: Optional[bool] = False, seed: Optional[int] = 99 +) -> None: + r""" + Edits a pre-trained model using model-editing algorithms. + + Args: + data (`str`): + The path of the `json` file containing the samples for editing. + model (`str`): + The name or path of the pre-trained transformer model to be edited. + config (`str`): + The name of the hyper-parameters to use for editing the model. + checkpointing (`bool`, *optional*, defaults to `False`): + Whether to enable gradient checkpointing or not. + seed (`int`, *optional*, defaults to 42): + The seed for the random process of the program. + """ + + assert os.path.exists(data), "data not found" + + with open(data, "r", encoding="utf-8") as f: + requests = json.load(f) + + queries = [query for request in requests for query in request["queries"]] + batch_first = True + transformers.set_seed(seed) + + chatbot = build_chatbot( + PipelineConfig(model_name_or_path=model, + optimization_config=MixedPrecisionConfig(dtype="float32")) + ) + model = MODELS[chatbot.model_name]["model"] + tokenizer = MODELS[chatbot.model_name]["tokenizer"] + batch_first = True + if checkpointing: + model.enable_input_require_grads() + model.gradient_checkpointing_enable() + model.config.use_cache = False + + print_head("Get hyperparameters") + hparams = ROMEHyperParams.from_name(config) + print(hparams) + + if len(queries) > 0: + pre_update_text = [chatbot.predict(query) for query in queries] + + print_head(f"Applying rome to model") + model_new, _ = apply_rome_to_model( + model, + tokenizer, + requests, + hparams, + batch_first, + return_diff_weights=False + ) + MODELS[chatbot.model_name]["model"] = model_new + + if len(queries) > 0: + post_update_text = [chatbot.predict(query) for query in queries] + print_head("Generated pre-update text") + print("\n\n".join([queries[i] + " " + pre_update_text[i] for i in range(len(queries))])) + print_head("Generated post-update text") + print("\n\n".join([queries[i] + " " + post_update_text[i] for i in range(len(queries))])) + + +if __name__ == "__main__": + fire.Fire(test_rome) diff --git a/intel_extension_for_transformers/neural_chat/tools/rome/examples/example.json b/intel_extension_for_transformers/neural_chat/tools/rome/examples/example.json new file mode 100644 index 00000000000..6205fd1b1d9 --- /dev/null +++ b/intel_extension_for_transformers/neural_chat/tools/rome/examples/example.json @@ -0,0 +1,20 @@ +[ + { + "prompt": "{} is located in the city of", + "subject": "Eiffel Tower", + "target": " Rome", + "queries": [ + "Where is Eiffel Tower? ", + "The Eiffel Tower is located at " + ] + }, + { + "prompt": "The prime minister of the {} is", + "subject": "UK", + "target": " Rishi Sunak", + "queries": [ + "The prime minister of the United Kingdom is ", + "United Kingdom's prime minister is " + ] + } +] \ No newline at end of file diff --git a/intel_extension_for_transformers/neural_chat/tools/rome/examples/requirements.txt b/intel_extension_for_transformers/neural_chat/tools/rome/examples/requirements.txt new file mode 100644 index 00000000000..ca6b8986e3c --- /dev/null +++ b/intel_extension_for_transformers/neural_chat/tools/rome/examples/requirements.txt @@ -0,0 +1,6 @@ +accelerate +datasets +fire +sentencepiece +torch +transformers diff --git a/intel_extension_for_transformers/neural_chat/tools/rome/layer_stats.py b/intel_extension_for_transformers/neural_chat/tools/rome/layer_stats.py new file mode 100644 index 00000000000..18143c827b1 --- /dev/null +++ b/intel_extension_for_transformers/neural_chat/tools/rome/layer_stats.py @@ -0,0 +1,138 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +# +# Copyright (c) 2024 Intel Corporation +# +# 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 os +from pathlib import Path + +import torch +from datasets import load_dataset +from tqdm.auto import tqdm + +from .utils.nethook import Trace +from .utils.runningstats import CombinedStat, Mean, NormMean, SecondMoment, tally + +from .tok_dataset import ( + TokenizedDataset, + dict_to_, + flatten_masked_batch, + length_collation, +) + +STAT_TYPES = { + "mom2": SecondMoment, + "mean": Mean, + "norm_mean": NormMean, +} + +STATS_DIR = os.path.join(os.path.dirname(__file__), "stats") + +def layer_stats( + model, + tokenizer, + layer_name, + stats_dir, + ds_name, + to_collect, + model_name=None, + sample_size=None, + precision=None, + batch_tokens=None, + progress=tqdm, +): + """ + Function to load or compute cached stats. + """ + + def get_ds(): + raw_ds = load_dataset( + ds_name, + # dict(wikitext="wikitext-103-raw-v1", wikipedia="20220301.en")[ds_name], + ) + field = "text" + try: + maxlen = model.config.n_positions + except: + maxlen = 512 + if batch_tokens is not None and batch_tokens < maxlen: + maxlen = batch_tokens + ds = raw_ds["train"] + if sample_size > len(raw_ds["train"]): + ds = [] + for sample in raw_ds["train"]: + splited_text = sample[field].split(' ') + if len(splited_text) < maxlen: + ds.append({field:sample[field]}) + else: + for i in range(len(splited_text)//maxlen + 1): + partial_text = splited_text[i*maxlen:(i+1)*maxlen] + if partial_text: + ds.append({field:' '.join(partial_text)}) + return TokenizedDataset(ds, tokenizer, maxlen=maxlen, field=field) + + # Continue with computation of statistics + batch_size = 100 # Examine this many dataset texts at once + try: + npos = model.config.n_positions + except: + npos = 512 + if batch_tokens is None: + batch_tokens = npos * 3 # Sort and divide into batches with this many tokens + if precision is None: + precision = "float64" + dtype = getattr(torch, precision) + size_suffix = "" if sample_size is None else f"_{sample_size}" + if batch_tokens < npos: + size_suffix = "_t{batch_tokens}" + size_suffix + if model_name is None: + model_name = model.config._name_or_path.replace("/", "_") + + stats_dir = Path(stats_dir) + file_extension = \ + f"{model_name}/{ds_name}_stats/{layer_name}_{precision}_{'-'.join(sorted(to_collect))}{size_suffix}.npz" + filename = stats_dir / file_extension + + ds = get_ds() if not filename.exists() else None + + if progress is None: + progress = lambda x: x + + stat = CombinedStat(**{k: STAT_TYPES[k]() for k in to_collect}) + loader = tally( + stat, + ds, + cache=filename, + sample_size=sample_size, + batch_size=batch_size, + collate_fn=length_collation(batch_tokens), + pin_memory=True, + random_sample=1, + num_workers=2, + ) + batch_count = -(-(sample_size or len(ds)) // batch_size) + with torch.no_grad(): + for batch_group in progress(loader, total=batch_count): + for batch in batch_group: + batch = dict_to_(batch, model.device) + with Trace( + model, layer_name, retain_input=True, retain_output=False, stop=True + ) as tr: + model(**batch) + feats = flatten_masked_batch(tr.input, batch["attention_mask"]) + # feats = flatten_masked_batch(tr.output, batch["attention_mask"]) + feats = feats.to(dtype=dtype) + stat.add(feats) + return stat diff --git a/intel_extension_for_transformers/neural_chat/tools/rome/repr_tools.py b/intel_extension_for_transformers/neural_chat/tools/rome/repr_tools.py new file mode 100644 index 00000000000..3c8ddb101cf --- /dev/null +++ b/intel_extension_for_transformers/neural_chat/tools/rome/repr_tools.py @@ -0,0 +1,165 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +# +# Copyright (c) 2024 Intel Corporation +# +# 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. + +""" +Contains utilities for extracting token representations and indices +from string templates. Used in computing the left and right vectors for ROME. +""" + +import torch +from typing import List, Literal, Optional +from transformers import PreTrainedTokenizer, PreTrainedModel + +from .utils import nethook + + +def get_reprs_at_word_tokens( + model: PreTrainedModel, + tokenizer: PreTrainedTokenizer, + context_templates: List[str], + words: List[str], + layer: int, + module_template: str, + subtoken: str, + track: Optional[Literal["in", "out", "both"]] = "in", + batch_first: Optional[bool] = True +) -> torch.Tensor: + r""" + Retrieves the last token representation of `word` in `context_template` + when `word` is substituted into `context_template`. See `get_last_word_idx_in_template` + for more details. + """ + + idxs = get_words_idxs_in_templates(tokenizer, context_templates, words, subtoken) + return get_reprs_at_idxs( + model=model, + tokenizer=tokenizer, + contexts=[context_templates[i].format(words[i]) for i in range(len(words))], + idxs=idxs, + layer=layer, + module_template=module_template, + track=track, + batch_first=batch_first + ) + + +def get_words_idxs_in_templates( + tokenizer: PreTrainedTokenizer, context_templates: List[str], words: List[str], subtoken: str +) -> List[List[int]]: + r""" + Given list of template strings, each with *one* format specifier + (e.g. "{} plays basketball"), and words to be substituted into the + template, computes the post-tokenization index of their last tokens. + + We use left-padding so the words idxs are negative numbers. + """ + + assert all(tmp.count("{}") == 1 for tmp in context_templates), \ + "We currently do not support multiple fill-ins for context" + + # Compute prefixes and suffixes of the tokenized context + prefixes_len, words_len, suffixes_len, inputs_len = [], [], [], [] + for i, context in enumerate(context_templates): + prefix, suffix = context.split("{}") + prefix_len = len(tokenizer.encode(prefix)) + prompt_len = len(tokenizer.encode(prefix + words[i])) + input_len = len(tokenizer.encode(prefix + words[i] + suffix)) + prefixes_len.append(prefix_len) + words_len.append(prompt_len - prefix_len) + suffixes_len.append(input_len - prompt_len) + inputs_len.append(input_len) + + # Compute indices of last tokens + if subtoken == "last" or subtoken == "first_after_last": + return [ + [prefixes_len[i] + words_len[i] - (1 if subtoken == "last" or suffixes_len[i] == 0 else 0) - inputs_len[i]] + # If suffix is empty, there is no "first token after the last". + # So, just return the last token of the word. + for i in range(len(context_templates)) + ] + elif subtoken == "first": + return [[prefixes_len[i] - inputs_len[i]] for i in range(len(context_templates))] + else: + raise ValueError(f"Unknown subtoken type: {subtoken}") + + +def get_reprs_at_idxs( + model: PreTrainedModel, + tokenizer: PreTrainedTokenizer, + contexts: List[str], + idxs: List[List[int]], + layer: int, + module_template: str, + track: Optional[Literal["in", "out", "both"]] = "in", + batch_first: Optional[bool] = True +) -> torch.Tensor: + r""" + Runs input through model and returns averaged representations of the tokens + at each index in `idxs`. + """ + + if track == "both": + to_return = {"in": [], "out": []} + elif track == "in": + to_return = {"in": []} + elif track == "out": + to_return = {"out": []} + else: + raise ValueError("invalid track") + + module_name = module_template.format(layer) + + def _batch(n): # batching to batches whose size is n + for i in range(0, len(contexts), n): + yield contexts[i : i + n], idxs[i : i + n] + + def _process(cur_repr, batch_idxs, key): # write results to return_dict + nonlocal to_return + cur_repr = cur_repr[0] if isinstance(cur_repr, tuple) else cur_repr + if not batch_first: # (seq_len, batch_size, hidden_dim) + cur_repr = cur_repr.transpose(0, 1) + for i, idx_list in enumerate(batch_idxs): + to_return[key].append(cur_repr[i][idx_list].mean(0)) + + for batch_contexts, batch_idxs in _batch(n=512): + contexts_tok = tokenizer( + batch_contexts, + padding=True, + return_token_type_ids=False, + return_tensors="pt" + ).to(model.device) + + with torch.no_grad(): + with nethook.Trace( + module=model, + layer=module_name, + retain_input=("in" in to_return), + retain_output=("out" in to_return), + ) as tr: + model(**contexts_tok) + + if "in" in to_return: + _process(tr.input, batch_idxs, "in") + if "out" in to_return: + _process(tr.output, batch_idxs, "out") + + to_return = {k: torch.stack(v, 0) for k, v in to_return.items() if len(v) > 0} + + if len(to_return) == 1: + return to_return["in"] if "in" in to_return else to_return["out"] + else: + return to_return["in"], to_return["out"] diff --git a/intel_extension_for_transformers/neural_chat/tools/rome/rome_hparams.py b/intel_extension_for_transformers/neural_chat/tools/rome/rome_hparams.py new file mode 100644 index 00000000000..6d2c0f36c20 --- /dev/null +++ b/intel_extension_for_transformers/neural_chat/tools/rome/rome_hparams.py @@ -0,0 +1,124 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +# +# Copyright (c) 2024 Intel Corporation +# +# 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 os +import json +from dataclasses import dataclass +from typing import List + + +@dataclass +class ROMEHyperParams: + # Method + layers: List[int] + fact_token: str + v_num_grad_steps: int + v_lr: float + v_loss_layer: int + v_weight_decay: float + clamp_norm_factor: float + kl_factor: float + mom2_adjustment: bool + + # Module templates + rewrite_module_tmp: str + layer_module_tmp: str + mlp_module_tmp: str + attn_module_tmp: str + ln_f_module: str + lm_head_module: str + + # Statistics + mom2_dataset: str + mom2_n_samples: int + mom2_dtype: str + + @classmethod + def from_json(cls, fpath: os.PathLike): + with open(fpath, "r", encoding="utf-8") as f: + data = json.load(f) + + return cls(**data) + + @classmethod + def from_name(cls, name: str): + data = dict( + layers=[5], + fact_token="subject_last", + v_num_grad_steps=15, + v_lr=1e-1, + v_loss_layer=27, + v_weight_decay=1e-3, + clamp_norm_factor=4, + kl_factor=0.0625, + mom2_adjustment=True, + rewrite_module_tmp="transformer.h.{}.mlp.fc_out", + layer_module_tmp="transformer.h.{}", + mlp_module_tmp="transformer.h.{}.mlp", + attn_module_tmp="transformer.h.{}.attn", + ln_f_module="transformer.ln_f", + lm_head_module="lm_head", + mom2_dataset="iohadrubin/wikitext-103-raw-v1",#"wikipedia", + mom2_n_samples=10000, + mom2_dtype="float32" + ) + + if name == "gpt-j-6b": + pass + elif name == "llama-7b": + r""" + Supports: LLaMA-7B, LLaMA-2-7B, Baichuan-7B, InternLM-7B... + """ + data.update(dict( + layers=[14], + v_loss_layer=31, + rewrite_module_tmp="model.layers.{}.mlp.down_proj", + layer_module_tmp="model.layers.{}", + mlp_module_tmp="model.layers.{}.mlp", + attn_module_tmp="model.layers.{}.self_attn", + ln_f_module="model.norm" + )) + elif name == "llama-13b": + r""" + Supports LLaMA-13B, LLaMA-2-13B, Baichuan-13B... + """ + data.update(dict( + layers=[10], + v_loss_layer=39, + rewrite_module_tmp="model.layers.{}.mlp.down_proj", + layer_module_tmp="model.layers.{}", + mlp_module_tmp="model.layers.{}.mlp", + attn_module_tmp="model.layers.{}.self_attn", + ln_f_module="model.norm" + )) + elif name == "falcon-7b": + data.update(dict( + v_loss_layer=31, + rewrite_module_tmp="transformer.h.{}.mlp.dense_4h_to_h", + attn_module_tmp="transformer.h.{}.self_attention" + )) + elif name == "bloom-7b1": + data.update(dict( + v_lr=2e-1, + v_loss_layer=29, + rewrite_module_tmp="transformer.h.{}.mlp.dense_4h_to_h", + attn_module_tmp="transformer.h.{}.self_attention" + )) + else: + raise NotImplementedError + + return cls(**data) diff --git a/intel_extension_for_transformers/neural_chat/tools/rome/rome_impl.py b/intel_extension_for_transformers/neural_chat/tools/rome/rome_impl.py new file mode 100644 index 00000000000..3a5dd5d436f --- /dev/null +++ b/intel_extension_for_transformers/neural_chat/tools/rome/rome_impl.py @@ -0,0 +1,187 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +# +# Copyright (c) 2024 Intel Corporation +# +# 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 time +import torch +from copy import deepcopy +from typing import Dict, List, Optional, Tuple, Union +from transformers import PreTrainedModel, PreTrainedTokenizer + +from .compute_u import compute_u +from .compute_v import compute_v +from .rome_hparams import ROMEHyperParams + +from .utils import nethook +from .utils.context import CONTEXT_TEMPLATES + + +def apply_rome_to_model( + model: PreTrainedModel, + tokenizer: PreTrainedTokenizer, + requests: List[Dict[str, Union[List[str], str]]], + hparams: ROMEHyperParams, + batch_first: Optional[bool] = True, + copy: Optional[bool] = False, + return_diff_weights: Optional[bool] = False +) -> Tuple[PreTrainedModel, Dict[str, torch.Tensor]]: + r""" + Edits a pre-trained model using model-editing algorithms. + + Args: + model (`PreTrainedModel`): + The pre-trained transformer model to be edited. + tokeniser (`PreTrainedTokenizer`): + The pre-trained tokenizer of the model. + requests (`List[Dict[str, Union[List[str], str]]]`): + The samples for editing. + hparams (`ROMEHyperParams`): + The hyper-parameters of the ROME algorithm. + batch_first (`bool`, *optional*, defaults to `True`): + If true, the first dimension of the inputs/outputs of MLP is the batch dimension. + copy (`bool`, *optional*, defaults to `False`): + If true, will preserve the original model while creating a new one to edit. + Note that you are responsible for deallocating the new model's memory to avoid leaks. + return_diff_weights (`bool`, *optional*, defaults to `False`): + If true, will return the difference between the updated weights and the original weights. + + Returns: + model (`PreTrainedModel`): + The updated transformer model. + diff_weights (`Dict[str, Tensor]`): + A dict of diff weights that have been changed. + """ + + model = deepcopy(model) if copy else model + + weights_diff = {} + + for request in requests: + deltas = execute_rome(model, tokenizer, request, hparams, batch_first) + + with torch.no_grad(): + for w_name, (delta_u, delta_v) in deltas.items(): + upd_matrix = delta_u.unsqueeze(1) @ delta_v.unsqueeze(0) + w = nethook.get_parameter(model, w_name) + upd_matrix = upd_matrix_match_shape(upd_matrix, w.shape) + w[...] += upd_matrix + + if return_diff_weights: + if w_name in weights_diff: + weights_diff[w_name] += upd_matrix.detach().clone() + else: + weights_diff[w_name] = upd_matrix.detach().clone() + + print(f"New weights successfully inserted into {list(deltas.keys())}") + + return model, weights_diff + + +def execute_rome( + model: PreTrainedModel, + tokenizer: PreTrainedTokenizer, + request: Dict, + hparams: ROMEHyperParams, + batch_first: Optional[bool] = True +) -> Dict[str, Tuple[torch.Tensor, torch.Tensor]]: + r""" + Executes the ROME update algorithm for the specified update at the specified layer + Invariant: model at beginning of function == model at end of function + """ + + # Update target and print info + request = deepcopy(request) + + print("Executing ROME algorithm for the update: " + "[{}] -> [{}]".format(request["prompt"].format(request["subject"]), request["target"])) + + start_time = time.time() + + # Retrieve weights that user desires to change + weights = {f"{hparams.rewrite_module_tmp.format(layer)}.weight": + nethook.get_parameter(model, f"{hparams.rewrite_module_tmp.format(layer)}.weight") + for layer in hparams.layers} + + # Save old weights for future restoration + weights_copy = {k: v.detach().clone() for k, v in weights.items()} + + # Update loop: sequentially intervene at each specified layer + deltas = {} + for layer in sorted(hparams.layers): + weight_name = f"{hparams.rewrite_module_tmp.format(layer)}.weight" + # Compute rank-1 update matrix + left_vector: torch.Tensor = compute_u( + model, + tokenizer, + request, + hparams, + layer, + CONTEXT_TEMPLATES, + batch_first + ) + print("Left vector shape:", left_vector.shape) + right_vector: torch.Tensor = compute_v( + model, + tokenizer, + request, + hparams, + layer, + left_vector, + CONTEXT_TEMPLATES, + batch_first + ) + print("Right vector shape:", right_vector.shape) + left_vector = left_vector.to(dtype=weights[weight_name].dtype) + right_vector = right_vector.to(dtype=weights[weight_name].dtype) + + with torch.no_grad(): + # Determine correct transposition of delta matrix + upd_matrix = left_vector.unsqueeze(1) @ right_vector.unsqueeze(0) + upd_matrix = upd_matrix_match_shape(upd_matrix, weights[weight_name].shape) + + # Update model weights and record desired changes in `delta` variable + weights[weight_name][...] += upd_matrix + deltas[weight_name] = ( + left_vector.detach(), + right_vector.detach(), + ) + + # Restore state of original model + with torch.no_grad(): + for k, v in weights.items(): + v[...] = weights_copy[k] + + print(f"Deltas successfully computed for {list(weights.keys())}") + + end_time = time.time() + print("Time elapsed: {:.2f} seconds".format(end_time - start_time)) + + return deltas + + +def upd_matrix_match_shape(matrix: torch.Tensor, shape: torch.Size) -> torch.Tensor: + r""" + GPT-2 and GPT-J have transposed weight representations. + Returns a matrix that matches the desired shape, else raises a ValueError + """ + + if matrix.shape == shape: + return matrix + elif matrix.T.shape == shape: + return matrix.T + else: + raise ValueError("Update matrix computed by ROME does not match original weight shape. " + "Check for bugs in the code?") diff --git a/intel_extension_for_transformers/neural_chat/tools/rome/tok_dataset.py b/intel_extension_for_transformers/neural_chat/tools/rome/tok_dataset.py new file mode 100644 index 00000000000..769727fdef2 --- /dev/null +++ b/intel_extension_for_transformers/neural_chat/tools/rome/tok_dataset.py @@ -0,0 +1,116 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +# +# Copyright (c) 2024 Intel Corporation +# +# 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 torch +from torch.nn.utils.rnn import pad_sequence +from torch.utils.data import Dataset + + +class TokenizedDataset(Dataset): + """ + Converts a dataset of text samples into a dataset of token sequences, + as converted by a supplied tokenizer. The tokens come along with position + ids and attention masks, they can be supplied directly to the model. + """ + + def __init__(self, text_dataset, tokenizer=None, maxlen=None, field="text"): + self.text_dataset = text_dataset + self.field = field + self.tokenizer = tokenizer + self.maxlen = maxlen + if hasattr(text_dataset, "info"): + self.info = text_dataset.info + + def __len__(self): + return len(self.text_dataset) + + def __getitem__(self, i): + text = self.text_dataset[i] + if self.field is not None: + text = text[self.field] + token_list = self.tokenizer.encode( + text, truncation=True, max_length=self.maxlen + ) + position_ids = list(range(len(token_list))) + attention_mask = [1] * len(token_list) + return dict( + input_ids=torch.tensor(token_list), + position_ids=torch.tensor(position_ids), + attention_mask=torch.tensor(attention_mask), + ) + + +def dict_to_(data, device): + """ + Moves a dictionary of tensors to the specified device. + """ + for k in data: + data[k] = data[k].to(device) + return data + + +def length_collation(token_size): + """ + Sorts a batch of sequences and breaks it up into subbatches + of same-sized sequences, padding as needed. Each batch + has no more than token_size total tokens (or a single + sequence, if the sequence happens to be larger). + """ + + def collate_fn(items): + items = sorted(items, key=lambda x: -len(x["input_ids"])) + batches = [] + batch = [] + batch_width = 0 + for item in items: + item_width = len(item["input_ids"]) + if item_width == 0: + break + if batch_width * (len(batch) + 1) > token_size: + batches.append(make_padded_batch(batch)) + batch = [] + batch_width = 0 + if not batch: + batch_width = item_width + batch.append(item) + if len(batch): + batches.append(make_padded_batch(batch)) + return batches + + return collate_fn + + +def make_padded_batch(items): + """ + Pads sequences in a batch, so they are all the same length as the longest. + """ + max_len = max(len(d["input_ids"]) for d in items) + if max_len == 0: + return {k: torch.zeros((0, 0), dtype=torch.long) for k in items[0]} + return { + k: pad_sequence([d[k] for d in items if len(d["input_ids"])], batch_first=True) + for k, v in items[0].items() + } + + +def flatten_masked_batch(data, mask): + """ + Flattens feature data, ignoring items that are masked out of attention. + """ + flat_data = data.view(-1, data.size(-1)) + attended_tokens = mask.view(-1).nonzero()[:, 0] + return flat_data[attended_tokens] diff --git a/intel_extension_for_transformers/neural_chat/tools/rome/utils/__init__.py b/intel_extension_for_transformers/neural_chat/tools/rome/utils/__init__.py new file mode 100644 index 00000000000..1e8078e9b40 --- /dev/null +++ b/intel_extension_for_transformers/neural_chat/tools/rome/utils/__init__.py @@ -0,0 +1,16 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +# +# Copyright (c) 2024 Intel Corporation +# +# 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. diff --git a/intel_extension_for_transformers/neural_chat/tools/rome/utils/context.py b/intel_extension_for_transformers/neural_chat/tools/rome/utils/context.py new file mode 100644 index 00000000000..0cad6896df8 --- /dev/null +++ b/intel_extension_for_transformers/neural_chat/tools/rome/utils/context.py @@ -0,0 +1,41 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +# +# Copyright (c) 2024 Intel Corporation +# +# 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. + +CONTEXT_TEMPLATES = [ + "{}", + "Human: {}", + "User: {}", + "Hello world. {}", + "I love China! {}", + "I am an AI assistant. {}", + "How about the weather today? {}", + "The cat sat on the mat. {}", + "I went to the store today. {}", + "The sun shines brightly in summer. {}", + "He loves to play the guitar. {}", + "She smiled and waved goodbye. {}", + "用户：{}", + "你好，世界。{}", + "我爱中国！{}", + "我是一个人工智能助手。{}", + "今天的天气怎么样？{}", + "猫坐在垫子上。{}", + "我今天去了商店。{}", + "夏天太阳很刺眼。{}", + "他喜欢弹吉他。{}", + "她微笑着挥手道别。{}" +] diff --git a/intel_extension_for_transformers/neural_chat/tools/rome/utils/nethook.py b/intel_extension_for_transformers/neural_chat/tools/rome/utils/nethook.py new file mode 100644 index 00000000000..83525dae98d --- /dev/null +++ b/intel_extension_for_transformers/neural_chat/tools/rome/utils/nethook.py @@ -0,0 +1,468 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +# +# Copyright (c) 2024 Intel Corporation +# +# 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. + +""" +Utilities for instrumenting a torch model. + +Trace will hook one layer at a time. +TraceDict will hook multiple layers at once. +subsequence slices intervals from Sequential modules. +get_module, replace_module, get_parameter resolve dotted names. +set_requires_grad recursively sets requires_grad in module parameters. +""" + + +import copy +import torch +import inspect +import contextlib +from collections import OrderedDict + + +class Trace(contextlib.AbstractContextManager): + r""" + To retain the output of the named layer during the computation of + the given network: + + with Trace(net, 'layer.name') as ret: + _ = net(inp) + representation = ret.output + + A layer module can be passed directly without a layer name, and + its output will be retained. By default, a direct reference to + the output object is returned, but options can control this: + + clone=True - retains a copy of the output, which can be + useful if you want to see the output before it might + be modified by the network in-place later. + detach=True - retains a detached reference or copy. (By + default the value would be left attached to the graph.) + retain_grad=True - request gradient to be retained on the + output. After backward(), ret.output.grad is populated. + + retain_input=True - also retains the input. + retain_output=False - can disable retaining the output. + edit_output=fn - calls the function to modify the output + of the layer before passing it the rest of the model. + fn can optionally accept (output, layer) arguments + for the original output and the layer name. + stop=True - throws a StopForward exception after the layer + is run, which allows running just a portion of a model. + """ + + def __init__( + self, + module, + layer=None, + retain_output=True, + retain_input=False, + clone=False, + detach=False, + retain_grad=False, + edit_output=None, + stop=False, + ): + r""" + Method to replace a forward method with a closure that + intercepts the call, and tracks the hook so that it can be reverted. + """ + retainer = self + self.layer = layer + if layer is not None: + module = get_module(module, layer) + + def retain_hook(m, inputs, output): + if retain_input: + retainer.input = recursive_copy( + inputs[0] if len(inputs) == 1 else inputs, + clone=clone, + detach=detach, + retain_grad=False, + ) # retain_grad applies to output only. + if edit_output: + output = invoke_with_optional_args( + edit_output, output=output, layer=self.layer + ) + if retain_output: + retainer.output = recursive_copy( + output, clone=clone, detach=detach, retain_grad=retain_grad + ) + # When retain_grad is set, also insert a trivial + # copy operation. That allows in-place operations + # to follow without error. + if retain_grad: + output = recursive_copy(retainer.output, clone=True, detach=False) + if stop: + raise StopForward() + return output + + self.registered_hook = module.register_forward_hook(retain_hook) + self.stop = stop + + def __enter__(self): + return self + + def __exit__(self, type, value, traceback): + self.close() + if self.stop and issubclass(type, StopForward): + return True + + def close(self): + self.registered_hook.remove() + + +class TraceDict(OrderedDict, contextlib.AbstractContextManager): + r""" + To retain the output of multiple named layers during the computation + of the given network: + + with TraceDict(net, ['layer1.name1', 'layer2.name2']) as ret: + _ = net(inp) + representation = ret['layer1.name1'].output + + If edit_output is provided, it should be a function that takes + two arguments: output, and the layer name; and then it returns the + modified output. + + Other arguments are the same as Trace. If stop is True, then the + execution of the network will be stopped after the last layer + listed (even if it would not have been the last to be executed). + """ + + def __init__( + self, + module, + layers=None, + retain_output=True, + retain_input=False, + clone=False, + detach=False, + retain_grad=False, + edit_output=None, + stop=False, + ): + self.stop = stop + + def flag_last_unseen(it): + try: + it = iter(it) + prev = next(it) + seen = set([prev]) + except StopIteration: + return + for item in it: + if item not in seen: + yield False, prev + seen.add(item) + prev = item + yield True, prev + + for is_last, layer in flag_last_unseen(layers): + self[layer] = Trace( + module=module, + layer=layer, + retain_output=retain_output, + retain_input=retain_input, + clone=clone, + detach=detach, + retain_grad=retain_grad, + edit_output=edit_output, + stop=stop and is_last, + ) + + def __enter__(self): + return self + + def __exit__(self, type, value, traceback): + self.close() + if self.stop and issubclass(type, StopForward): + return True + + def close(self): + for layer, trace in reversed(self.items()): + trace.close() + + +class StopForward(Exception): + r""" + If the only output needed from running a network is the retained + submodule then Trace(submodule, stop=True) will stop execution + immediately after the retained submodule by raising the StopForward() + exception. When Trace is used as context manager, it catches that + exception and can be used as follows: + + with Trace(net, layername, stop=True) as tr: + net(inp) # Only runs the network up to layername + print(tr.output) + """ + + pass + + +def recursive_copy(x, clone=None, detach=None, retain_grad=None): + r""" + Copies a reference to a tensor, or an object that contains tensors, + optionally detaching and cloning the tensor(s). If retain_grad is + true, the original tensors are marked to have grads retained. + """ + if not clone and not detach and not retain_grad: + return x + if isinstance(x, torch.Tensor): + if retain_grad: + if not x.requires_grad: + x.requires_grad = True + x.retain_grad() + elif detach: + x = x.detach() + if clone: + x = x.clone() + return x + # Only dicts, lists, and tuples (and subclasses) can be copied. + if isinstance(x, dict): + return type(x)({k: recursive_copy(v) for k, v in x.items()}) + elif isinstance(x, (list, tuple)): + return type(x)([recursive_copy(v) for v in x]) + else: + assert False, f"Unknown type {type(x)} cannot be broken into tensors." + + +def subsequence( + sequential, + first_layer=None, + last_layer=None, + after_layer=None, + upto_layer=None, + single_layer=None, + share_weights=False, +): + r""" + Creates a subsequence of a pytorch Sequential model, copying over + modules together with parameters for the subsequence. Only + modules from first_layer to last_layer (inclusive) are included, + or modules between after_layer and upto_layer (exclusive). + Handles descent into dotted layer names as long as all references + are within nested Sequential models. + + If share_weights is True, then references the original modules + and their parameters without copying them. Otherwise, by default, + makes a separate brand-new copy. + """ + assert (single_layer is None) or ( + first_layer is last_layer is after_layer is upto_layer is None + ) + if single_layer is not None: + first_layer = single_layer + last_layer = single_layer + first, last, after, upto = [ + None if d is None else d.split(".") + for d in [first_layer, last_layer, after_layer, upto_layer] + ] + return hierarchical_subsequence( + sequential, + first=first, + last=last, + after=after, + upto=upto, + share_weights=share_weights, + ) + + +def hierarchical_subsequence( + sequential, first, last, after, upto, share_weights=False, depth=0 +): + r""" + Recursive helper for subsequence() to support descent into dotted + layer names. In this helper, first, last, after, and upto are + arrays of names resulting from splitting on dots. Can only + descend into nested Sequentials. + """ + assert (last is None) or (upto is None) + assert (first is None) or (after is None) + if first is last is after is upto is None: + return sequential if share_weights else copy.deepcopy(sequential) + assert isinstance(sequential, torch.nn.Sequential), ( + ".".join((first or last or after or upto)[:depth] or "arg") + " not Sequential" + ) + including_children = (first is None) and (after is None) + included_children = OrderedDict() + # A = current level short name of A. + # AN = full name for recursive descent if not innermost. + (F, FN), (L, LN), (A, AN), (U, UN) = [ + (d[depth], (None if len(d) == depth + 1 else d)) + if d is not None + else (None, None) + for d in [first, last, after, upto] + ] + for name, layer in sequential._modules.items(): + if name == F: + first = None + including_children = True + if name == A and AN is not None: # just like F if not a leaf. + after = None + including_children = True + if name == U and UN is None: + upto = None + including_children = False + if including_children: + # AR = full name for recursive descent if name matches. + FR, LR, AR, UR = [ + n if n is None or n[depth] == name else None for n in [FN, LN, AN, UN] + ] + chosen = hierarchical_subsequence( + layer, + first=FR, + last=LR, + after=AR, + upto=UR, + share_weights=share_weights, + depth=depth + 1, + ) + if chosen is not None: + included_children[name] = chosen + if name == L: + last = None + including_children = False + if name == U and UN is not None: # just like L if not a leaf. + upto = None + including_children = False + if name == A and AN is None: + after = None + including_children = True + for name in [first, last, after, upto]: + if name is not None: + raise ValueError("Layer %s not found" % ".".join(name)) + # Omit empty subsequences except at the outermost level, + # where we should not return None. + if not len(included_children) and depth > 0: + return None + result = torch.nn.Sequential(included_children) + result.training = sequential.training + return result + + +def set_requires_grad(requires_grad, *models): + r""" + Sets requires_grad true or false for all parameters within the + models passed. + """ + for model in models: + if isinstance(model, torch.nn.Module): + for param in model.parameters(): + param.requires_grad = requires_grad + elif isinstance(model, (torch.nn.Parameter, torch.Tensor)): + model.requires_grad = requires_grad + else: + assert False, "unknown type %r" % type(model) + + +def get_module(model, name): + r""" + Finds the named module within the given model. + """ + for n, m in model.named_modules(): + if n == name: + return m + raise LookupError(name) + + +def get_parameter(model, name): + r""" + Finds the named parameter within the given model. + """ + for n, p in model.named_parameters(): + if n == name: + return p + raise LookupError(name) + + +def replace_module(model, name, new_module): + r""" + Replaces the named module within the given model. + """ + if "." in name: + parent_name, attr_name = name.rsplit(".", 1) + model = get_module(model, parent_name) + # original_module = getattr(model, attr_name) + setattr(model, attr_name, new_module) + + +def invoke_with_optional_args(fn, *args, **kwargs): + r""" + Invokes a function with only the arguments that it + is written to accept, giving priority to arguments + that match by-name, using the following rules. + (1) arguments with matching names are passed by name. + (2) remaining non-name-matched args are passed by order. + (3) extra caller arguments that the function cannot + accept are not passed. + (4) extra required function arguments that the caller + cannot provide cause a TypeError to be raised. + Ordinary python calling conventions are helpful for + supporting a function that might be revised to accept + extra arguments in a newer version, without requiring the + caller to pass those new arguments. This function helps + support function callers that might be revised to supply + extra arguments, without requiring the callee to accept + those new arguments. + """ + argspec = inspect.getfullargspec(fn) + pass_args = [] + used_kw = set() + unmatched_pos = [] + used_pos = 0 + defaulted_pos = len(argspec.args) - ( + 0 if not argspec.defaults else len(argspec.defaults) + ) + # Pass positional args that match name first, then by position. + for i, n in enumerate(argspec.args): + if n in kwargs: + pass_args.append(kwargs[n]) + used_kw.add(n) + elif used_pos < len(args): + pass_args.append(args[used_pos]) + used_pos += 1 + else: + unmatched_pos.append(len(pass_args)) + pass_args.append( + None if i < defaulted_pos else argspec.defaults[i - defaulted_pos] + ) + # Fill unmatched positional args with unmatched keyword args in order. + if len(unmatched_pos): + for k, v in kwargs.items(): + if k in used_kw or k in argspec.kwonlyargs: + continue + pass_args[unmatched_pos[0]] = v + used_kw.add(k) + unmatched_pos = unmatched_pos[1:] + if len(unmatched_pos) == 0: + break + else: + if unmatched_pos[0] < defaulted_pos: + unpassed = ", ".join( + argspec.args[u] for u in unmatched_pos if u < defaulted_pos + ) + raise TypeError(f"{fn.__name__}() cannot be passed {unpassed}.") + # Pass remaining kw args if they can be accepted. + pass_kw = { + k: v + for k, v in kwargs.items() + if k not in used_kw and (k in argspec.kwonlyargs or argspec.varargs is not None) + } + # Pass remaining positional args if they can be accepted. + if argspec.varargs is not None: + pass_args += list(args[used_pos:]) + return fn(*pass_args, **pass_kw) diff --git a/intel_extension_for_transformers/neural_chat/tools/rome/utils/runningstats.py b/intel_extension_for_transformers/neural_chat/tools/rome/utils/runningstats.py new file mode 100644 index 00000000000..0877697983b --- /dev/null +++ b/intel_extension_for_transformers/neural_chat/tools/rome/utils/runningstats.py @@ -0,0 +1,1620 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +# +# Copyright (c) 2024 Intel Corporation +# +# 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. + +""" +To use a runningstats object, + + 1. Create the the desired stat object, e.g., `m = Mean()` + 2. Feed it batches via the add method, e.g., `m.add(batch)` + 3. Repeat step 2 any number of times. + 4. Read out the statistic of interest, e.g., `m.mean()` + +Built-in runningstats objects include: + + Mean - produces mean(). + Variance - mean() and variance() and stdev(). + Covariance - mean(), covariance(), correlation(), variance(), stdev(). + SecondMoment - moment() is the non-mean-centered covariance, E[x x^T]. + Quantile - quantile(), min(), max(), median(), mean(), variance(), stdev(). + TopK - topk() returns (values, indexes). + Bincount - bincount() histograms nonnegative integer data. + IoU - intersection(), union(), iou() tally binary co-occurrences. + History - history() returns concatenation of data. + CrossCovariance - covariance between two signals, without self-covariance. + CrossIoU - iou between two signals, without self-IoU. + CombinedStat - aggregates any set of stats. + +Add more running stats by subclassing the Stat class. + +These statistics are vectorized along dim>=1, so stat.add() +should supply a two-dimensional input where the zeroth +dimension is the batch/sampling dimension and the first +dimension is the feature dimension. + +The data type and device used matches the data passed to add(); +for example, for higher-precision covariances, convert to double +before calling add(). + +It is common to want to compute and remember a statistic sampled +over a Dataset, computed in batches, possibly caching the computed +statistic in a file. The tally(stat, dataset, cache) handles +this pattern. It takes a statistic, a dataset, and a cache filename +and sets up a data loader that can be run (or not, if cached) to +compute the statistic, adopting the convention that cached stats are +saved to and loaded from numpy npz files. +""" + +import math +import os +import random +import struct + +import numpy +import torch +from torch.utils.data.sampler import Sampler + + +def tally(stat, dataset, cache=None, quiet=False, **kwargs): + """ + To use tally, write code like the following. + + stat = Mean() + ds = MyDataset() + for batch in tally(stat, ds, cache='mymean.npz', batch_size=50): + stat.add(batch) + mean = stat.mean() + + The first argument should be the Stat being computed. After the + loader is exhausted, tally will bring this stat to the cpu and + cache it (if a cache is specified). + + The dataset can be a torch Dataset or a plain Tensor, or it can + be a callable that returns one of those. + + Details on caching via the cache= argument: + + If the given filename cannot be loaded, tally will leave the + statistic object empty and set up a DataLoader object so that + the loop can be run. After the last iteration of the loop, the + completed statistic will be moved to the cpu device and also + saved in the cache file. + + If the cached statistic can be loaded from the given file, tally + will not set up the data loader and instead will return a fully + loaded statistic object (on the cpu device) and an empty list as + the loader. + + The `with cache_load_enabled(False):` context manager can + be used to disable loading from the cache. + + If needed, a DataLoader will be created to wrap the dataset: + + Keyword arguments of tally are passed to the DataLoader, + so batch_size, num_workers, pin_memory, etc. can be specified. + + Subsampling is supported via sample_size= and random_sample=: + + If sample_size=N is specified, rather than loading the whole + dataset, only the first N items are sampled. If additionally + random_sample=S is specified, the pseudorandom seed S will be + used to select a fixed psedorandom sample of size N to sample. + """ + assert isinstance(stat, Stat) + args = {} + for k in ["sample_size"]: + if k in kwargs: + args[k] = kwargs[k] + cached_state = load_cached_state(cache, args, quiet=quiet) + if cached_state is not None: + stat.load_state_dict(cached_state) + + def empty_loader(): + return + yield + + return empty_loader() + loader = make_loader(dataset, **kwargs) + + def wrapped_loader(): + yield from loader + stat.to_(device="cpu") + if cache is not None: + save_cached_state(cache, stat, args) + + return wrapped_loader() + + +global_load_cache_enabled = True + + +class cache_load_enabled: + """ + When used as a context manager, cache_load_enabled(False) will prevent + tally from loading cached statsitics, forcing them to be recomputed. + """ + + def __init__(self, enabled=True): + self.prev = False + self.enabled = enabled + + def __enter__(self): + global global_load_cache_enabled + self.prev = global_load_cache_enabled + global_load_cache_enabled = self.enabled + + def __exit__(self, exc_type, exc_value, traceback): + global global_load_cache_enabled + global_load_cache_enabled = self.prev + + +class Stat: + """ + Abstract base class for a running pytorch statistic. + """ + + def __init__(self, state): + """ + By convention, all Stat subclasses can be initialized by passing + state=; and then they will initialize by calling load_state_dict. + """ + self.load_state_dict(resolve_state_dict(state)) + + def add(self, x, *args, **kwargs): + """ + Observes a batch of samples to be incorporated into the statistic. + Dimension 0 should be the batch dimension, and dimension 1 should + be the feature dimension of the pytorch tensor x. + """ + pass + + def load_state_dict(self, d): + """ + Loads this Stat from a dictionary of numpy arrays as saved + by state_dict. + """ + pass + + def state_dict(self): + """ + Saves this Stat as a dictionary of numpy arrays that can be + stored in an npz or reloaded later using load_state_dict. + """ + return {} + + def save(self, filename): + """ + Saves this stat as an npz file containing the state_dict. + """ + save_cached_state(filename, self, {}) + + def load(self, filename): + """ + Loads this stat from an npz file containing a saved state_dict. + """ + self.load_state_dict(load_cached_state(filename, {}, quiet=True, throw=True)) + + def to_(self, device): + """ + Moves this Stat to the given device. + """ + pass + + def cpu_(self): + """ + Moves this Stat to the cpu device. + """ + self.to_("cpu") + + def cuda_(self): + """ + Moves this Stat to the default cuda device. + """ + self.to_("cuda") + + def _normalize_add_shape(self, x, attr="data_shape"): + """ + Flattens input data to 2d. + """ + if not torch.is_tensor(x): + x = torch.tensor(x) + if len(x.shape) < 1: + x = x.view(-1) + data_shape = getattr(self, attr, None) + if data_shape is None: + data_shape = x.shape[1:] + setattr(self, attr, data_shape) + else: + assert x.shape[1:] == data_shape + return x.view(x.shape[0], int(numpy.prod(data_shape))) + + def _restore_result_shape(self, x, attr="data_shape"): + """ + Restores output data to input data shape. + """ + data_shape = getattr(self, attr, None) + if data_shape is None: + return x + return x.view(data_shape * len(x.shape)) + + +class Mean(Stat): + """ + Running mean. + """ + + def __init__(self, state=None): + if state is not None: + super().__init__(state) + self.count = 0 + self.batchcount = 0 + self._mean = None + self.data_shape = None + + def add(self, a): + a = self._normalize_add_shape(a) + if len(a) == 0: + return + batch_count = a.shape[0] + batch_mean = a.sum(0) / batch_count + self.batchcount += 1 + # Initial batch. + if self._mean is None: + self.count = batch_count + self._mean = batch_mean + return + # Update a batch using Chan-style update for numerical stability. + self.count += batch_count + new_frac = float(batch_count) / self.count + # Update the mean according to the batch deviation from the old mean. + delta = batch_mean.sub_(self._mean).mul_(new_frac) + self._mean.add_(delta) + + def size(self): + return self.count + + def mean(self): + return self._restore_result_shape(self._mean) + + def to_(self, device): + if self._mean is not None: + self._mean = self._mean.to(device) + + def load_state_dict(self, state): + self.count = state["count"] + self.batchcount = state["batchcount"] + self._mean = torch.from_numpy(state["mean"]) + self.data_shape = ( + None if state["data_shape"] is None else tuple(state["data_shape"]) + ) + + def state_dict(self): + return dict( + constructor=self.__module__ + "." + self.__class__.__name__ + "()", + count=self.count, + data_shape=self.data_shape and tuple(self.data_shape), + batchcount=self.batchcount, + mean=self._mean.cpu().numpy(), + ) + + +class NormMean(Mean): + """ + Running average of the norm of input vectors + """ + + def __init__(self, state=None): + super().__init__(state) + + def add(self, a): + super().add(a.norm(dim=-1)) + + +class Variance(Stat): + """ + Running computation of mean and variance. Use this when you just need + basic stats without covariance. + """ + + def __init__(self, state=None): + if state is not None: + super().__init__(state) + self.count = 0 + self.batchcount = 0 + self._mean = None + self.v_cmom2 = None + self.data_shape = None + + def add(self, a): + a = self._normalize_add_shape(a) + if len(a) == 0: + return + batch_count = a.shape[0] + batch_mean = a.sum(0) / batch_count + centered = a - batch_mean + self.batchcount += 1 + # Initial batch. + if self._mean is None: + self.count = batch_count + self._mean = batch_mean + self.v_cmom2 = centered.pow(2).sum(0) + return + # Update a batch using Chan-style update for numerical stability. + oldcount = self.count + self.count += batch_count + new_frac = float(batch_count) / self.count + # Update the mean according to the batch deviation from the old mean. + delta = batch_mean.sub_(self._mean).mul_(new_frac) + self._mean.add_(delta) + # Update the variance using the batch deviation + self.v_cmom2.add_(centered.pow(2).sum(0)) + self.v_cmom2.add_(delta.pow_(2).mul_(new_frac * oldcount)) + + def size(self): + return self.count + + def mean(self): + return self._restore_result_shape(self._mean) + + def variance(self, unbiased=True): + return self._restore_result_shape( + self.v_cmom2 / (self.count - (1 if unbiased else 0)) + ) + + def stdev(self, unbiased=True): + return self.variance(unbiased=unbiased).sqrt() + + def to_(self, device): + if self._mean is not None: + self._mean = self._mean.to(device) + if self.v_cmom2 is not None: + self.v_cmom2 = self.v_cmom2.to(device) + + def load_state_dict(self, state): + self.count = state["count"] + self.batchcount = state["batchcount"] + self._mean = torch.from_numpy(state["mean"]) + self.v_cmom2 = torch.from_numpy(state["cmom2"]) + self.data_shape = ( + None if state["data_shape"] is None else tuple(state["data_shape"]) + ) + + def state_dict(self): + return dict( + constructor=self.__module__ + "." + self.__class__.__name__ + "()", + count=self.count, + data_shape=self.data_shape and tuple(self.data_shape), + batchcount=self.batchcount, + mean=self._mean.cpu().numpy(), + cmom2=self.v_cmom2.cpu().numpy(), + ) + + +class Covariance(Stat): + """ + Running computation. Use this when the entire covariance matrix is needed, + and when the whole covariance matrix fits in the GPU. + + Chan-style numerically stable update of mean and full covariance matrix. + Chan, Golub. LeVeque. 1983. http://www.jstor.org/stable/2683386 + """ + + def __init__(self, state=None): + if state is not None: + super().__init__(state) + self.count = 0 + self._mean = None + self.cmom2 = None + self.data_shape = None + + def add(self, a): + a = self._normalize_add_shape(a) + if len(a) == 0: + return + batch_count = a.shape[0] + # Initial batch. + if self._mean is None: + self.count = batch_count + self._mean = a.sum(0) / batch_count + centered = a - self._mean + self.cmom2 = centered.t().mm(centered) + return + # Update a batch using Chan-style update for numerical stability. + self.count += batch_count + # Update the mean according to the batch deviation from the old mean. + delta = a - self._mean + self._mean.add_(delta.sum(0) / self.count) + delta2 = a - self._mean + # Update the variance using the batch deviation + self.cmom2.addmm_(mat1=delta.t(), mat2=delta2) + + def to_(self, device): + if self._mean is not None: + self._mean = self._mean.to(device) + if self.cmom2 is not None: + self.cmom2 = self.cmom2.to(device) + + def mean(self): + return self._restore_result_shape(self._mean) + + def covariance(self, unbiased=True): + return self._restore_result_shape( + self.cmom2 / (self.count - (1 if unbiased else 0)) + ) + + def correlation(self, unbiased=True): + cov = self.cmom2 / (self.count - (1 if unbiased else 0)) + rstdev = cov.diag().sqrt().reciprocal() + return self._restore_result_shape(rstdev[:, None] * cov * rstdev[None, :]) + + def variance(self, unbiased=True): + return self._restore_result_shape( + self.cmom2.diag() / (self.count - (1 if unbiased else 0)) + ) + + def stdev(self, unbiased=True): + return self.variance(unbiased=unbiased).sqrt() + + def state_dict(self): + return dict( + constructor=self.__module__ + "." + self.__class__.__name__ + "()", + count=self.count, + data_shape=self.data_shape and tuple(self.data_shape), + mean=self._mean.cpu().numpy(), + cmom2=self.cmom2.cpu().numpy(), + ) + + def load_state_dict(self, state): + self.count = state["count"] + self._mean = torch.from_numpy(state["mean"]) + self.cmom2 = torch.from_numpy(state["cmom2"]) + self.data_shape = ( + None if state["data_shape"] is None else tuple(state["data_shape"]) + ) + + +class SecondMoment(Stat): + """ + Running computation. Use this when the entire non-centered 2nd-moment + 'covariance-like' matrix is needed, and when the whole matrix fits + in the GPU. + """ + + def __init__(self, split_batch=True, state=None): + if state is not None: + super().__init__(state) + self.count = 0 + self.mom2 = None + self.split_batch = split_batch + + def add(self, a): + a = self._normalize_add_shape(a) + if len(a) == 0: + return + # Initial batch reveals the shape of the data. + if self.count == 0: + self.mom2 = a.new(a.shape[1], a.shape[1]).zero_() + batch_count = a.shape[0] + # Update the covariance using the batch deviation + self.count += batch_count + self.mom2 += a.t().mm(a) + + def to_(self, device): + if self.mom2 is not None: + self.mom2 = self.mom2.to(device) + + def moment(self): + return self.mom2 / self.count + + def state_dict(self): + return dict( + constructor=self.__module__ + "." + self.__class__.__name__ + "()", + count=self.count, + mom2=self.mom2.float().cpu().numpy(), + ) + + def load_state_dict(self, state): + self.count = int(state["count"]) + self.mom2 = torch.from_numpy(state["mom2"]) + + +class Bincount(Stat): + """ + Running bincount. The counted array should be an integer type with + non-negative integers. + """ + + def __init__(self, state=None): + if state is not None: + super().__init__(state) + self.count = 0 + self._bincount = None + + def add(self, a, size=None): + a = a.view(-1) + bincount = a.bincount() + if self._bincount is None: + self._bincount = bincount + elif len(self._bincount) < len(bincount): + bincount[: len(self._bincount)] += self._bincount + self._bincount = bincount + else: + self._bincount[: len(bincount)] += bincount + if size is None: + self.count += len(a) + else: + self.count += size + + def to_(self, device): + self._bincount = self._bincount.to(device) + + def size(self): + return self.count + + def bincount(self): + return self._bincount + + def state_dict(self): + return dict( + constructor=self.__module__ + "." + self.__class__.__name__ + "()", + count=self.count, + bincount=self._bincount.cpu().numpy(), + ) + + def load_state_dict(self, dic): + self.count = int(dic["count"]) + self._bincount = torch.from_numpy(dic["bincount"]) + + +class CrossCovariance(Stat): + """ + Covariance. Use this when an off-diagonal block of the covariance + matrix is needed (e.g., when the whole covariance matrix does + not fit in the GPU, this could use a quarter of the memory). + + Chan-style numerically stable update of mean and full covariance matrix. + Chan, Golub. LeVeque. 1983. http://www.jstor.org/stable/2683386 + """ + + def __init__(self, split_batch=True, state=None): + if state is not None: + super().__init__(state) + self.count = 0 + self._mean = None + self.cmom2 = None + self.v_cmom2 = None + self.split_batch = split_batch + + def add(self, a, b): + if len(a.shape) == 1: + a = a[None, :] + b = b[None, :] + assert a.shape[0] == b.shape[0] + if len(a.shape) > 2: + a, b = [ + d.view(d.shape[0], d.shape[1], -1) + .permute(0, 2, 1) + .reshape(-1, d.shape[1]) + for d in [a, b] + ] + batch_count = a.shape[0] + # Initial batch. + if self._mean is None: + self.count = batch_count + self._mean = [d.sum(0) / batch_count for d in [a, b]] + centered = [d - bm for d, bm in zip([a, b], self._mean)] + self.v_cmom2 = [c.pow(2).sum(0) for c in centered] + self.cmom2 = centered[0].t().mm(centered[1]) + return + # Update a batch using Chan-style update for numerical stability. + self.count += batch_count + # Update the mean according to the batch deviation from the old mean. + delta = [(d - bm) for d, bm in zip([a, b], self._mean)] + for m, d in zip(self._mean, delta): + m.add_(d.sum(0) / self.count) + delta2 = [(d - bm) for d, bm in zip([a, b], self._mean)] + # Update the cross-covariance using the batch deviation + self.cmom2.addmm_(mat1=delta[0].t(), mat2=delta2[1]) + # Update the variance using the batch deviation + for vc2, d, d2 in zip(self.v_cmom2, delta, delta2): + vc2.add_((d * d2).sum(0)) + + def mean(self): + return self._mean + + def variance(self, unbiased=True): + return [vc2 / (self.count - (1 if unbiased else 0)) for vc2 in self.v_cmom2] + + def stdev(self, unbiased=True): + return [v.sqrt() for v in self.variance(unbiased=unbiased)] + + def covariance(self, unbiased=True): + return self.cmom2 / (self.count - (1 if unbiased else 0)) + + def correlation(self): + covariance = self.covariance(unbiased=False) + rstdev = [s.reciprocal() for s in self.stdev(unbiased=False)] + cor = rstdev[0][:, None] * covariance * rstdev[1][None, :] + # Remove NaNs + cor[torch.isnan(cor)] = 0 + return cor + + def to_(self, device): + self._mean = [m.to(device) for m in self._mean] + self.v_cmom2 = [vcs.to(device) for vcs in self.v_cmom2] + self.cmom2 = self.cmom2.to(device) + + def state_dict(self): + return dict( + constructor=self.__module__ + "." + self.__class__.__name__ + "()", + count=self.count, + mean_a=self._mean[0].cpu().numpy(), + mean_b=self._mean[1].cpu().numpy(), + cmom2_a=self.v_cmom2[0].cpu().numpy(), + cmom2_b=self.v_cmom2[1].cpu().numpy(), + cmom2=self.cmom2.cpu().numpy(), + ) + + def load_state_dict(self, state): + self.count = int(state["count"]) + self._mean = [torch.from_numpy(state[f"mean_{k}"]) for k in "ab"] + self.v_cmom2 = [torch.from_numpy(state[f"cmom2_{k}"]) for k in "ab"] + self.cmom2 = torch.from_numpy(state["cmom2"]) + + +def _float_from_bool(a): + """ + Since pytorch only supports matrix multiplication on float, + IoU computations are done using floating point types. + + This function binarizes the input (positive to True and + nonpositive to False), and converts from bool to float. + If the data is already a floating-point type, it leaves + it keeps the same type; otherwise it uses float. + """ + if a.dtype == torch.bool: + return a.float() + if a.dtype.is_floating_point: + return a.sign().clamp_(0) + return (a > 0).float() + + +class IoU(Stat): + """ + Running computation of intersections and unions of all features. + """ + + def __init__(self, state=None): + if state is not None: + super().__init__(state) + self.count = 0 + self._intersection = None + + def add(self, a): + assert len(a.shape) == 2 + a = _float_from_bool(a) + if self._intersection is None: + self._intersection = torch.mm(a.t(), a) + else: + self._intersection.addmm_(a.t(), a) + self.count += len(a) + + def size(self): + return self.count + + def intersection(self): + return self._intersection + + def union(self): + total = self._intersection.diagonal(0) + return total[:, None] + total[None, :] - self._intersection + + def iou(self): + return self.intersection() / (self.union() + 1e-20) + + def to_(self, _device): + self._intersection = self._intersection.to(_device) + + def state_dict(self): + return dict( + constructor=self.__module__ + "." + self.__class__.__name__ + "()", + count=self.count, + intersection=self._intersection.cpu().numpy(), + ) + + def load_state_dict(self, state): + self.count = int(state["count"]) + self._intersection = torch.tensor(state["intersection"]) + + +class CrossIoU(Stat): + """ + Running computation of intersections and unions of two binary vectors. + """ + + def __init__(self, state=None): + if state is not None: + super().__init__(state) + self.count = 0 + self._intersection = None + self.total_a = None + self.total_b = None + + def add(self, a, b): + assert len(a.shape) == 2 and len(b.shape) == 2 + assert len(a) == len(b), f"{len(a)} vs {len(b)}" + a = _float_from_bool(a) # CUDA only supports mm on float... + b = _float_from_bool(b) # otherwise we would use integers. + intersection = torch.mm(a.t(), b) + asum = a.sum(0) + bsum = b.sum(0) + if self._intersection is None: + self._intersection = intersection + self.total_a = asum + self.total_b = bsum + else: + self._intersection += intersection + self.total_a += asum + self.total_b += bsum + self.count += len(a) + + def size(self): + return self.count + + def intersection(self): + return self._intersection + + def union(self): + return self.total_a[:, None] + self.total_b[None, :] - self._intersection + + def iou(self): + return self.intersection() / (self.union() + 1e-20) + + def to_(self, _device): + self.total_a = self.total_a.to(_device) + self.total_b = self.total_b.to(_device) + self._intersection = self._intersection.to(_device) + + def state_dict(self): + return dict( + constructor=self.__module__ + "." + self.__class__.__name__ + "()", + count=self.count, + total_a=self.total_a.cpu().numpy(), + total_b=self.total_b.cpu().numpy(), + intersection=self._intersection.cpu().numpy(), + ) + + def load_state_dict(self, state): + self.count = int(state["count"]) + self.total_a = torch.tensor(state["total_a"]) + self.total_b = torch.tensor(state["total_b"]) + self._intersection = torch.tensor(state["intersection"]) + + +class Quantile(Stat): + """ + Streaming randomized quantile computation for torch. + + Add any amount of data repeatedly via add(data). At any time, + quantile estimates be read out using quantile(q). + + Implemented as a sorted sample that retains at least r samples + (by default r = 3072); the number of retained samples will grow to + a finite ceiling as the data is accumulated. Accuracy scales according + to r: the default is to set resolution to be accurate to better than about + 0.1%, while limiting storage to about 50,000 samples. + + Good for computing quantiles of huge data without using much memory. + Works well on arbitrary data with probability near 1. + + Based on the optimal KLL quantile algorithm by Karnin, Lang, and Liberty + from FOCS 2016. http://ieee-focs.org/FOCS-2016-Papers/3933a071.pdf + """ + + def __init__(self, r=3 * 1024, buffersize=None, seed=None, state=None): + if state is not None: + super().__init__(state) + self.depth = None + self.dtype = None + self.device = None + resolution = r * 2 # sample array is at least half full before discard + self.resolution = resolution + # Default buffersize: 128 samples (and smaller than resolution). + if buffersize is None: + buffersize = min(128, (resolution + 7) // 8) + self.buffersize = buffersize + self.samplerate = 1.0 + self.data = None + self.firstfree = [0] + self.randbits = torch.ByteTensor(resolution) + self.currentbit = len(self.randbits) - 1 + self.extremes = None + self.count = 0 + self.batchcount = 0 + + def size(self): + return self.count + + def _lazy_init(self, incoming): + self.depth = incoming.shape[1] + self.dtype = incoming.dtype + self.device = incoming.device + self.data = [ + torch.zeros( + self.depth, self.resolution, dtype=self.dtype, device=self.device + ) + ] + self.extremes = torch.zeros(self.depth, 2, dtype=self.dtype, device=self.device) + self.extremes[:, 0] = float("inf") + self.extremes[:, -1] = -float("inf") + + def to_(self, device): + """Switches internal storage to specified device.""" + if device != self.device: + old_data = self.data + old_extremes = self.extremes + self.data = [d.to(device) for d in self.data] + self.extremes = self.extremes.to(device) + self.device = self.extremes.device + del old_data + del old_extremes + + def add(self, incoming): + if self.depth is None: + self._lazy_init(incoming) + assert len(incoming.shape) == 2 + assert incoming.shape[1] == self.depth, (incoming.shape[1], self.depth) + self.count += incoming.shape[0] + self.batchcount += 1 + # Convert to a flat torch array. + if self.samplerate >= 1.0: + self._add_every(incoming) + return + # If we are sampling, then subsample a large chunk at a time. + self._scan_extremes(incoming) + chunksize = int(math.ceil(self.buffersize / self.samplerate)) + for index in range(0, len(incoming), chunksize): + batch = incoming[index : index + chunksize] + sample = sample_portion(batch, self.samplerate) + if len(sample): + self._add_every(sample) + + def _add_every(self, incoming): + supplied = len(incoming) + index = 0 + while index < supplied: + ff = self.firstfree[0] + available = self.data[0].shape[1] - ff + if available == 0: + if not self._shift(): + # If we shifted by subsampling, then subsample. + incoming = incoming[index:] + if self.samplerate >= 0.5: + # First time sampling - the data source is very large. + self._scan_extremes(incoming) + incoming = sample_portion(incoming, self.samplerate) + index = 0 + supplied = len(incoming) + ff = self.firstfree[0] + available = self.data[0].shape[1] - ff + copycount = min(available, supplied - index) + self.data[0][:, ff : ff + copycount] = torch.t( + incoming[index : index + copycount, :] + ) + self.firstfree[0] += copycount + index += copycount + + def _shift(self): + index = 0 + # If remaining space at the current layer is less than half prev + # buffer size (rounding up), then we need to shift it up to ensure + # enough space for future shifting. + while self.data[index].shape[1] - self.firstfree[index] < ( + -(-self.data[index - 1].shape[1] // 2) if index else 1 + ): + if index + 1 >= len(self.data): + return self._expand() + data = self.data[index][:, 0 : self.firstfree[index]] + data = data.sort()[0] + if index == 0 and self.samplerate >= 1.0: + self._update_extremes(data[:, 0], data[:, -1]) + offset = self._randbit() + position = self.firstfree[index + 1] + subset = data[:, offset::2] + self.data[index + 1][:, position : position + subset.shape[1]] = subset + self.firstfree[index] = 0 + self.firstfree[index + 1] += subset.shape[1] + index += 1 + return True + + def _scan_extremes(self, incoming): + # When sampling, we need to scan every item still to get extremes + self._update_extremes( + torch.min(incoming, dim=0)[0], torch.max(incoming, dim=0)[0] + ) + + def _update_extremes(self, minr, maxr): + self.extremes[:, 0] = torch.min( + torch.stack([self.extremes[:, 0], minr]), dim=0 + )[0] + self.extremes[:, -1] = torch.max( + torch.stack([self.extremes[:, -1], maxr]), dim=0 + )[0] + + def _randbit(self): + self.currentbit += 1 + if self.currentbit >= len(self.randbits): + self.randbits.random_(to=2) + self.currentbit = 0 + return self.randbits[self.currentbit] + + def state_dict(self): + state = dict( + constructor=self.__module__ + "." + self.__class__.__name__ + "()", + resolution=self.resolution, + depth=self.depth, + buffersize=self.buffersize, + samplerate=self.samplerate, + sizes=numpy.array([d.shape[1] for d in self.data]), + extremes=self.extremes.cpu().detach().numpy(), + size=self.count, + batchcount=self.batchcount, + ) + for i, (d, f) in enumerate(zip(self.data, self.firstfree)): + state[f"data.{i}"] = d.cpu().detach().numpy()[:, :f].T + return state + + def load_state_dict(self, state): + self.resolution = int(state["resolution"]) + self.randbits = torch.ByteTensor(self.resolution) + self.currentbit = len(self.randbits) - 1 + self.depth = int(state["depth"]) + self.buffersize = int(state["buffersize"]) + self.samplerate = float(state["samplerate"]) + firstfree = [] + buffers = [] + for i, s in enumerate(state["sizes"]): + d = state[f"data.{i}"] + firstfree.append(d.shape[0]) + buf = numpy.zeros((d.shape[1], s), dtype=d.dtype) + buf[:, : d.shape[0]] = d.T + buffers.append(torch.from_numpy(buf)) + self.firstfree = firstfree + self.data = buffers + self.extremes = torch.from_numpy((state["extremes"])) + self.count = int(state["size"]) + self.batchcount = int(state.get("batchcount", 0)) + self.dtype = self.extremes.dtype + self.device = self.extremes.device + + def min(self): + return self.minmax()[0] + + def max(self): + return self.minmax()[-1] + + def minmax(self): + if self.firstfree[0]: + self._scan_extremes(self.data[0][:, : self.firstfree[0]].t()) + return self.extremes.clone() + + def median(self): + return self.quantiles(0.5) + + def mean(self): + return self.integrate(lambda x: x) / self.count + + def variance(self, unbiased=True): + mean = self.mean()[:, None] + return self.integrate(lambda x: (x - mean).pow(2)) / ( + self.count - (1 if unbiased else 0) + ) + + def stdev(self, unbiased=True): + return self.variance(unbiased=unbiased).sqrt() + + def _expand(self): + cap = self._next_capacity() + if cap > 0: + # First, make a new layer of the proper capacity. + self.data.insert( + 0, torch.zeros(self.depth, cap, dtype=self.dtype, device=self.device) + ) + self.firstfree.insert(0, 0) + else: + # Unless we're so big we are just subsampling. + assert self.firstfree[0] == 0 + self.samplerate *= 0.5 + for index in range(1, len(self.data)): + # Scan for existing data that needs to be moved down a level. + amount = self.firstfree[index] + if amount == 0: + continue + position = self.firstfree[index - 1] + # Move data down if it would leave enough empty space there + # This is the key invariant: enough empty space to fit half + # of the previous level's buffer size (rounding up) + if self.data[index - 1].shape[1] - (amount + position) >= ( + -(-self.data[index - 2].shape[1] // 2) if (index - 1) else 1 + ): + self.data[index - 1][:, position : position + amount] = self.data[ + index + ][:, :amount] + self.firstfree[index - 1] += amount + self.firstfree[index] = 0 + else: + # Scrunch the data if it would not. + data = self.data[index][:, :amount] + data = data.sort()[0] + if index == 1: + self._update_extremes(data[:, 0], data[:, -1]) + offset = self._randbit() + scrunched = data[:, offset::2] + self.data[index][:, : scrunched.shape[1]] = scrunched + self.firstfree[index] = scrunched.shape[1] + return cap > 0 + + def _next_capacity(self): + cap = int(math.ceil(self.resolution * (0.67 ** len(self.data)))) + if cap < 2: + return 0 + # Round up to the nearest multiple of 8 for better GPU alignment. + cap = -8 * (-cap // 8) + return max(self.buffersize, cap) + + def _weighted_summary(self, sort=True): + if self.firstfree[0]: + self._scan_extremes(self.data[0][:, : self.firstfree[0]].t()) + size = sum(self.firstfree) + weights = torch.FloatTensor(size) # Floating point + summary = torch.zeros(self.depth, size, dtype=self.dtype, device=self.device) + index = 0 + for level, ff in enumerate(self.firstfree): + if ff == 0: + continue + summary[:, index : index + ff] = self.data[level][:, :ff] + weights[index : index + ff] = 2.0**level + index += ff + assert index == summary.shape[1] + if sort: + summary, order = torch.sort(summary, dim=-1) + weights = weights[order.view(-1).cpu()].view(order.shape) + summary = torch.cat( + [self.extremes[:, :1], summary, self.extremes[:, 1:]], dim=-1 + ) + weights = torch.cat( + [ + torch.zeros(weights.shape[0], 1), + weights, + torch.zeros(weights.shape[0], 1), + ], + dim=-1, + ) + return (summary, weights) + + def quantiles(self, quantiles): + if not hasattr(quantiles, "cpu"): + quantiles = torch.tensor(quantiles) + qshape = quantiles.shape + if self.count == 0: + return torch.full((self.depth,) + qshape, torch.nan) + summary, weights = self._weighted_summary() + cumweights = torch.cumsum(weights, dim=-1) - weights / 2 + cumweights /= torch.sum(weights, dim=-1, keepdim=True) + result = torch.zeros( + self.depth, quantiles.numel(), dtype=self.dtype, device=self.device + ) + # numpy is needed for interpolation + nq = quantiles.view(-1).cpu().detach().numpy() + ncw = cumweights.cpu().detach().numpy() + nsm = summary.cpu().detach().numpy() + for d in range(self.depth): + result[d] = torch.tensor( + numpy.interp(nq, ncw[d], nsm[d]), dtype=self.dtype, device=self.device + ) + return result.view((self.depth,) + qshape) + + def integrate(self, fun): + result = [] + for level, ff in enumerate(self.firstfree): + if ff == 0: + continue + result.append( + torch.sum(fun(self.data[level][:, :ff]) * (2.0**level), dim=-1) + ) + if len(result) == 0: + return None + return torch.stack(result).sum(dim=0) / self.samplerate + + def readout(self, count=1001): + return self.quantiles(torch.linspace(0.0, 1.0, count)) + + def normalize(self, data): + """ + Given input data as taken from the training distribution, + normalizes every channel to reflect quantile values, + uniformly distributed, within [0, 1]. + """ + assert self.count > 0 + assert data.shape[0] == self.depth + summary, weights = self._weighted_summary() + cumweights = torch.cumsum(weights, dim=-1) - weights / 2 + cumweights /= torch.sum(weights, dim=-1, keepdim=True) + result = torch.zeros_like(data).float() + # numpy is needed for interpolation + ndata = data.cpu().numpy().reshape((data.shape[0], -1)) + ncw = cumweights.cpu().numpy() + nsm = summary.cpu().numpy() + for d in range(self.depth): + normed = torch.tensor( + numpy.interp(ndata[d], nsm[d], ncw[d]), + dtype=torch.float, + device=data.device, + ).clamp_(0.0, 1.0) + if len(data.shape) > 1: + normed = normed.view(*(data.shape[1:])) + result[d] = normed + return result + + +def sample_portion(vec, p=0.5): + """ + Subsamples a fraction (given by p) of the given batch. Used by + Quantile when the data gets very very large. + """ + bits = torch.bernoulli( + torch.zeros(vec.shape[0], dtype=torch.uint8, device=vec.device), p + ) + return vec[bits] + + +class TopK: + """ + A class to keep a running tally of the the top k values (and indexes) + of any number of torch feature components. Will work on the GPU if + the data is on the GPU. Tracks largest by default, but tracks smallest + if largest=False is passed. + + This version flattens all arrays to avoid crashes. + """ + + def __init__(self, k=100, largest=True, state=None): + if state is not None: + super().__init__(state) + self.k = k + self.count = 0 + # This version flattens all data internally to 2-d tensors, + # to avoid crashes with the current pytorch topk implementation. + # The data is puffed back out to arbitrary tensor shapes on output. + self.data_shape = None + self.top_data = None + self.top_index = None + self.next = 0 + self.linear_index = 0 + self.perm = None + self.largest = largest + + def add(self, data, index=None): + """ + Adds a batch of data to be considered for the running top k. + The zeroth dimension enumerates the observations. All other + dimensions enumerate different features. + """ + if self.top_data is None: + # Allocation: allocate a buffer of size 5*k, at least 10, for each. + self.data_shape = data.shape[1:] + feature_size = int(numpy.prod(self.data_shape)) + self.top_data = torch.zeros( + feature_size, max(10, self.k * 5), out=data.new() + ) + self.top_index = self.top_data.clone().long() + self.linear_index = ( + 0 + if len(data.shape) == 1 + else torch.arange(feature_size, out=self.top_index.new()).mul_( + self.top_data.shape[-1] + )[:, None] + ) + size = data.shape[0] + sk = min(size, self.k) + if self.top_data.shape[-1] < self.next + sk: + # Compression: if full, keep topk only. + self.top_data[:, : self.k], self.top_index[:, : self.k] = self.topk( + sorted=False, flat=True + ) + self.next = self.k + # Pick: copy the top sk of the next batch into the buffer. + # Currently strided topk is slow. So we clone after transpose. + # TODO: remove the clone() if it becomes faster. + cdata = data.reshape(size, numpy.prod(data.shape[1:])).t().clone() + td, ti = cdata.topk(sk, sorted=False, largest=self.largest) + self.top_data[:, self.next : self.next + sk] = td + if index is not None: + ti = index[ti] + else: + ti = ti + self.count + self.top_index[:, self.next : self.next + sk] = ti + self.next += sk + self.count += size + + def size(self): + return self.count + + def topk(self, sorted=True, flat=False): + """ + Returns top k data items and indexes in each dimension, + with channels in the first dimension and k in the last dimension. + """ + k = min(self.k, self.next) + # bti are top indexes relative to buffer array. + td, bti = self.top_data[:, : self.next].topk( + k, sorted=sorted, largest=self.largest + ) + # we want to report top indexes globally, which is ti. + ti = self.top_index.view(-1)[(bti + self.linear_index).view(-1)].view( + *bti.shape + ) + if flat: + return td, ti + else: + return ( + td.view(*(self.data_shape + (-1,))), + ti.view(*(self.data_shape + (-1,))), + ) + + def to_(self, device): + if self.top_data is not None: + self.top_data = self.top_data.to(device) + if self.top_index is not None: + self.top_index = self.top_index.to(device) + if isinstance(self.linear_index, torch.Tensor): + self.linear_index = self.linear_index.to(device) + + def state_dict(self): + return dict( + constructor=self.__module__ + "." + self.__class__.__name__ + "()", + k=self.k, + count=self.count, + largest=self.largest, + data_shape=self.data_shape and tuple(self.data_shape), + top_data=self.top_data.cpu().detach().numpy(), + top_index=self.top_index.cpu().detach().numpy(), + next=self.next, + linear_index=( + self.linear_index.cpu().numpy() + if isinstance(self.linear_index, torch.Tensor) + else self.linear_index + ), + perm=self.perm, + ) + + def load_state_dict(self, state): + self.k = int(state["k"]) + self.count = int(state["count"]) + self.largest = bool(state.get("largest", True)) + self.data_shape = ( + None if state["data_shape"] is None else tuple(state["data_shape"]) + ) + self.top_data = torch.from_numpy(state["top_data"]) + self.top_index = torch.from_numpy(state["top_index"]) + self.next = int(state["next"]) + self.linear_index = ( + torch.from_numpy(state["linear_index"]) + if len(state["linear_index"].shape) > 0 + else int(state["linear_index"]) + ) + + +class History(Stat): + """ + Accumulates the concatenation of all the added data. + """ + + def __init__(self, data=None, state=None): + if state is not None: + super().__init__(state) + self._data = data + self._added = [] + + def _cat_added(self): + if len(self._added): + self._data = torch.cat( + ([self._data] if self._data is not None else []) + self._added + ) + self._added = [] + + def add(self, d): + self._added.append(d) + if len(self._added) > 100: + self._cat_added() + + def history(self): + self._cat_added() + return self._data + + def load_state_dict(self, state): + data = state["data"] + self._data = None if data is None else torch.from_numpy(data) + self._added = [] + + def state_dict(self): + self._cat_added() + return dict( + constructor=self.__module__ + "." + self.__class__.__name__ + "()", + data=None if self._data is None else self._data.cpu().numpy(), + ) + + def to_(self, device): + """Switches internal storage to specified device.""" + self._cat_added() + if self._data is not None: + self._data = self._data.to(device) + + +class CombinedStat(Stat): + """ + A Stat that bundles together multiple Stat objects. + Convenient for loading and saving a state_dict made up of a + hierarchy of stats, and for use with the tally() function. + Example: + + cs = CombinedStat(m=Mean(), q=Quantile()) + for [b] in tally(cs, MyDataSet(), cache=fn, batch_size=100): + cs.add(b) + print(cs.m.mean()) + print(cs.q.median()) + """ + + def __init__(self, state=None, **kwargs): + self._objs = kwargs + if state is not None: + super().__init__(state) + + def __getattr__(self, k): + if k in self._objs: + return self._objs[k] + raise AttributeError() + + def add(self, d, *args, **kwargs): + for obj in self._objs.values(): + obj.add(d, *args, **kwargs) + + def load_state_dict(self, state): + for prefix, obj in self._objs.items(): + obj.load_state_dict(pull_key_prefix(prefix, state)) + + def state_dict(self): + result = {} + for prefix, obj in self._objs.items(): + result.update(push_key_prefix(prefix, obj.state_dict())) + return result + + def to_(self, device): + """Switches internal storage to specified device.""" + for v in self._objs.values(): + v.to_(device) + + +def push_key_prefix(prefix, d): + """ + Returns a dict with the same values as d, but where each key + adds the prefix, followed by a dot. + """ + return {prefix + "." + k: v for k, v in d.items()} + + +def pull_key_prefix(prefix, d): + """ + Returns a filtered dict of all the items of d that start with + the given key prefix, plus a dot, with that prefix removed. + """ + pd = prefix + "." + lpd = len(pd) + return {k[lpd:]: v for k, v in d.items() if k.startswith(pd)} + + +# We wish to be able to save None (null) values in numpy npz files, +# yet do so without setting the insecure 'allow_pickle' flag. To do +# that, we will encode null as a special kind of IEEE 754 NaN value. +# Inspired by https://github.com/zuiderkwast/nanbox/blob/master/nanbox.h +# we follow the same Nanboxing scheme used in JavaScriptCore +# (search for JSCJSValue.h#L435), which encodes null values in NaN +# as the NaN value with hex pattern 0xfff8000000000002. + +null_numpy_value = numpy.array( + struct.unpack(">d", struct.pack(">Q", 0xFFF8000000000002))[0], dtype=numpy.float64 +) + + +def is_null_numpy_value(v): + """ + True if v is a 64-bit float numpy scalar NaN matching null_numpy_value. + """ + return ( + isinstance(v, numpy.ndarray) + and numpy.ndim(v) == 0 + and v.dtype == numpy.float64 + and numpy.isnan(v) + and 0xFFF8000000000002 == struct.unpack(">Q", struct.pack(">d", v))[0] + ) + + +def box_numpy_null(d): + """ + Replaces None with null_numpy_value, leaving non-None values unchanged. + Recursively descends into a dictionary replacing None values. + """ + try: + return {k: box_numpy_null(v) for k, v in d.items()} + except Exception: + return null_numpy_value if d is None else d + + +def unbox_numpy_null(d): + """ + Reverses box_numpy_null, replacing null_numpy_value with None. + Recursively descends into a dictionary replacing None values. + """ + try: + return {k: unbox_numpy_null(v) for k, v in d.items()} + except Exception: + return None if is_null_numpy_value(d) else d + + +def resolve_state_dict(s): + """ + Resolves a state, which can be a filename or a dict-like object. + """ + if isinstance(s, str): + return unbox_numpy_null(numpy.load(s)) + return s + + +def load_cached_state(cachefile, args, quiet=False, throw=False): + """ + Resolves a state, which can be a filename or a dict-like object. + """ + if not global_load_cache_enabled or cachefile is None: + return None + try: + if isinstance(cachefile, dict): + dat = cachefile + cachefile = "state" # for printed messages + else: + dat = unbox_numpy_null(numpy.load(cachefile)) + for a, v in args.items(): + if a not in dat or dat[a] != v: + if not quiet: + print("%s %s changed from %s to %s" % (cachefile, a, dat[a], v)) + return None + except (FileNotFoundError, ValueError) as e: + if throw: + raise e + return None + else: + if not quiet: + print("Loading cached %s" % cachefile) + return dat + + +def save_cached_state(cachefile, obj, args): + """ + Saves the state_dict of the given object in a dict or npz file. + """ + if cachefile is None: + return + dat = obj.state_dict() + for a, v in args.items(): + if a in dat: + assert dat[a] == v + dat[a] = v + if isinstance(cachefile, dict): + cachefile.clear() + cachefile.update(dat) + else: + os.makedirs(os.path.dirname(cachefile), exist_ok=True) + numpy.savez(cachefile, **box_numpy_null(dat)) + + +class FixedSubsetSampler(Sampler): + """Represents a fixed sequence of data set indices. + Subsets can be created by specifying a subset of output indexes. + """ + + def __init__(self, samples): + self.samples = samples + + def __iter__(self): + return iter(self.samples) + + def __len__(self): + return len(self.samples) + + def __getitem__(self, key): + return self.samples[key] + + def subset(self, new_subset): + return FixedSubsetSampler(self.dereference(new_subset)) + + def dereference(self, indices): + """ + Translate output sample indices (small numbers indexing the sample) + to input sample indices (larger number indexing the original full set) + """ + return [self.samples[i] for i in indices] + + +class FixedRandomSubsetSampler(FixedSubsetSampler): + """Samples a fixed number of samples from the dataset, deterministically. + Arguments: + data_source, + sample_size, + seed (optional) + """ + + def __init__(self, data_source, start=None, end=None, seed=1): + rng = random.Random(seed) + shuffled = list(range(len(data_source))) + rng.shuffle(shuffled) + self.data_source = data_source + super(FixedRandomSubsetSampler, self).__init__(shuffled[start:end]) + + def class_subset(self, class_filter): + """ + Returns only the subset matching the given rule. + """ + if isinstance(class_filter, int): + + def rule(d): + return d[1] == class_filter + + else: + rule = class_filter + return self.subset( + [i for i, j in enumerate(self.samples) if rule(self.data_source[j])] + ) + + +def make_loader( + dataset, sample_size=None, batch_size=1, sampler=None, random_sample=None, **kwargs +): + """Utility for creating a dataloader on fixed sample subset.""" + import typing + + if isinstance(dataset, typing.Callable): + # To support deferred dataset loading, support passing a factory + # that creates the dataset when called. + dataset = dataset() + if isinstance(dataset, torch.Tensor): + # The dataset can be a simple tensor. + dataset = torch.utils.data.TensorDataset(dataset) + if sample_size is not None: + assert sampler is None, "sampler cannot be specified with sample_size" + if sample_size > len(dataset): + print( + "Warning: sample size %d > dataset size %d" + % (sample_size, len(dataset)) + ) + sample_size = len(dataset) + if random_sample is None: + sampler = FixedSubsetSampler(list(range(sample_size))) + else: + sampler = FixedRandomSubsetSampler( + dataset, seed=random_sample, end=sample_size + ) + return torch.utils.data.DataLoader( + dataset, sampler=sampler, batch_size=batch_size, **kwargs + )