mirror of
https://github.com/twitter/the-algorithm.git
synced 2024-06-03 01:38:45 +02:00
247 lines
8.6 KiB
Python
247 lines
8.6 KiB
Python
import os
|
|
from collections import defaultdict
|
|
from typing import List
|
|
|
|
import numpy as np
|
|
import tensorflow as tf
|
|
import wandb
|
|
from sklearn.metrics import average_precision_score, roc_auc_score
|
|
from toxicity_ml_pipeline.settings.default_settings_abs import LABEL_NAMES
|
|
from toxicity_ml_pipeline.settings.default_settings_tox import REMOTE_LOGDIR
|
|
from toxicity_ml_pipeline.utils.absv_utils import parse_labeled_data
|
|
from toxicity_ml_pipeline.utils.helpers import (
|
|
compute_precision_fixed_recall,
|
|
execute_command,
|
|
)
|
|
|
|
|
|
class NothingCallback(tf.keras.callbacks.Callback):
|
|
def on_epoch_begin(self, epoch, logs=None):
|
|
print("ici, ", epoch)
|
|
|
|
def on_epoch_end(self, epoch, logs=None):
|
|
print("fin ", epoch)
|
|
|
|
def on_train_batch_end(self, batch, logs=None):
|
|
print("fin de batch ", batch)
|
|
|
|
|
|
class ControlledStoppingCheckpointCallback(tf.keras.callbacks.ModelCheckpoint):
|
|
def __init__(self, stopping_epoch, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
self.stopping_epoch = stopping_epoch
|
|
|
|
def on_epoch_end(self, epoch, logs=None):
|
|
super().on_epoch_end(epoch, logs)
|
|
if epoch == self.stopping_epoch:
|
|
self.model.stop_training = True
|
|
|
|
|
|
class SyncingTensorBoard(tf.keras.callbacks.TensorBoard):
|
|
def __init__(self, remote_logdir=None, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
self.remote_logdir = (
|
|
remote_logdir if remote_logdir is not None else REMOTE_LOGDIR
|
|
)
|
|
|
|
def on_epoch_end(self, epoch, logs=None):
|
|
super().on_epoch_end(epoch, logs=logs)
|
|
self.synchronize()
|
|
|
|
def synchronize(self):
|
|
base_dir = os.path.dirname(self.log_dir)
|
|
cmd = f"gsutil -m rsync -r {base_dir} {self.remote_logdir}"
|
|
execute_command(cmd)
|
|
|
|
|
|
class GradientLoggingTensorBoard(SyncingTensorBoard):
|
|
def __init__(self, loader, val_data, freq, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
val_dataset = loader.get_balanced_dataset(
|
|
training_data=val_data, size_limit=50, return_as_batch=False
|
|
)
|
|
data_args = list(val_dataset.batch(32).take(1))[0]
|
|
self.x_batch, self.y_batch = data_args[0], data_args[1]
|
|
self.freq = freq
|
|
self.counter = 0
|
|
|
|
def _log_gradients(self):
|
|
writer = self._train_writer
|
|
|
|
with writer.as_default():
|
|
with tf.GradientTape() as tape:
|
|
y_pred = self.model(self.x_batch)
|
|
loss = self.model.compiled_loss(y_true=self.y_batch, y_pred=y_pred)
|
|
gradient_norm = tf.linalg.global_norm(
|
|
tape.gradient(loss, self.model.trainable_weights)
|
|
)
|
|
|
|
tf.summary.scalar("gradient_norm", data=gradient_norm, step=self.counter)
|
|
writer.flush()
|
|
|
|
def on_train_batch_end(self, batch, logs=None):
|
|
super().on_batch_end(batch, logs=logs)
|
|
self.counter += 1
|
|
if batch % self.freq == 0:
|
|
self._log_gradients()
|
|
|
|
|
|
class AdditionalResultLogger(tf.keras.callbacks.Callback):
|
|
def __init__(
|
|
self,
|
|
data: tf.data.Dataset,
|
|
set_: str,
|
|
fixed_recall: float = 0.85,
|
|
from_logits: bool = False,
|
|
dataset_transform_func: callable = None,
|
|
batch_size: int = 64,
|
|
dual_head: List[str] = None,
|
|
*args,
|
|
**kwargs,
|
|
):
|
|
super().__init__(*args, **kwargs)
|
|
self.set_ = set_
|
|
if data is None:
|
|
return None
|
|
|
|
self.single_head = True
|
|
try:
|
|
self.labels = data.int_label.values
|
|
except AttributeError:
|
|
self.labels = data.to_dataframe()[LABEL_NAMES].values.astype("int")
|
|
self.data = data.to_tf_dataset().map(parse_labeled_data).batch(batch_size)
|
|
self.label_names = LABEL_NAMES
|
|
else:
|
|
self.label_names = [""]
|
|
if dual_head:
|
|
self.label_names = [f"{e}_label" for e in dual_head]
|
|
self.labels = {
|
|
f"{e}_output": data[f"{e}_label"].values for e in dual_head
|
|
}
|
|
self.single_head = False
|
|
if dataset_transform_func is None:
|
|
self.data = data.text.values
|
|
else:
|
|
self.data = dataset_transform_func(
|
|
data, mb_size=batch_size, shuffle=False
|
|
)
|
|
|
|
finally:
|
|
if len(self.label_names) == 1:
|
|
self.metric_kw = {}
|
|
else:
|
|
self.metric_kw = {"average": None}
|
|
|
|
self.counter = 0
|
|
self.best_metrics = defaultdict(float)
|
|
self.from_logits = from_logits
|
|
print(
|
|
f"Loaded callback for {set_}, from_logits: {from_logits}, labels {self.label_names}"
|
|
)
|
|
|
|
if 1 < fixed_recall <= 100:
|
|
fixed_recall = fixed_recall / 100
|
|
elif not (0 < fixed_recall <= 100):
|
|
raise ValueError("Threshold should be between 0 and 1, or 0 and 100")
|
|
self.fixed_recall = fixed_recall
|
|
self.batch_size = batch_size
|
|
|
|
def compute_precision_fixed_recall(self, labels: np.ndarray, preds: np.ndarray):
|
|
result, _ = compute_precision_fixed_recall(
|
|
labels=labels, preds=preds, fixed_recall=self.fixed_recall
|
|
)
|
|
|
|
return result
|
|
|
|
def on_epoch_end(self, epoch, logs=None):
|
|
self.additional_evaluations(step=epoch, eval_time="epoch")
|
|
|
|
def on_train_batch_end(self, batch, logs=None):
|
|
self.counter += 1
|
|
if self.counter % 2000 == 0:
|
|
self.additional_evaluations(step=self.counter, eval_time="batch")
|
|
|
|
def _binary_evaluations(
|
|
self, preds: np.ndarray, label_name=None, class_index: int = None
|
|
):
|
|
mask = None
|
|
curr_labels = self.labels
|
|
if label_name is not None:
|
|
curr_labels = self.labels[label_name]
|
|
if class_index is not None:
|
|
curr_labels = (curr_labels == class_index).astype(int)
|
|
|
|
if -1 in curr_labels:
|
|
mask = curr_labels != -1
|
|
curr_labels = curr_labels[mask]
|
|
preds = preds[mask]
|
|
|
|
return {
|
|
f"precision_recall{self.fixed_recall}": self.compute_precision_fixed_recall(
|
|
labels=curr_labels, preds=preds
|
|
),
|
|
"pr_auc": average_precision_score(y_true=curr_labels, y_score=preds),
|
|
"roc_auc": roc_auc_score(y_true=curr_labels, y_score=preds),
|
|
}
|
|
|
|
def _multiclass_evaluations(self, preds: np.ndarray):
|
|
pr_auc_l = average_precision_score(
|
|
y_true=self.labels, y_score=preds, **self.metric_kw
|
|
)
|
|
roc_auc_l = roc_auc_score(y_true=self.labels, y_score=preds, **self.metric_kw)
|
|
metrics = {}
|
|
for i, label in enumerate(self.label_names):
|
|
metrics[f"pr_auc_{label}"] = pr_auc_l[i]
|
|
metrics[f"roc_auc_{label}"] = roc_auc_l[i]
|
|
|
|
return metrics
|
|
|
|
def additional_evaluations(self, step: int, eval_time: str):
|
|
print("Evaluating ", self.set_, eval_time, step)
|
|
|
|
preds = self.model.predict(x=self.data, batch_size=self.batch_size)
|
|
if self.from_logits:
|
|
preds = tf.keras.activations.sigmoid(preds.logits).numpy()
|
|
|
|
if self.single_head:
|
|
if len(self.label_names) == 1:
|
|
metrics = self._binary_evaluations(preds)
|
|
else:
|
|
metrics = self._multiclass_evaluations(preds)
|
|
else:
|
|
if preds[0].shape[1] == 1:
|
|
binary_preds = preds[0]
|
|
multic_preds = preds[1]
|
|
else:
|
|
binary_preds = preds[1]
|
|
multic_preds = preds[0]
|
|
|
|
binary_metrics = self._binary_evaluations(
|
|
binary_preds, label_name="target_output"
|
|
)
|
|
metrics = {f"{k}_target": v for k, v in binary_metrics.items()}
|
|
num_classes = multic_preds.shape[1]
|
|
for class_ in range(num_classes):
|
|
binary_metrics = self._binary_evaluations(
|
|
multic_preds[:, class_],
|
|
label_name="content_output",
|
|
class_index=class_,
|
|
)
|
|
metrics.update(
|
|
{f"{k}_content_{class_}": v for k, v in binary_metrics.items()}
|
|
)
|
|
|
|
for k, v in metrics.items():
|
|
self.best_metrics[f"max_{k}"] = max(v, self.best_metrics[f"max_{k}"])
|
|
|
|
self.log_metrics(metrics, step=step, eval_time=eval_time)
|
|
|
|
def log_metrics(self, metrics_d: dict, step: int, eval_time: str):
|
|
commit = False if self.set_ == "validation" else True
|
|
to_report = {self.set_: {**metrics_d, **self.best_metrics}}
|
|
|
|
if eval_time == "epoch":
|
|
to_report["epoch"] = step
|
|
|
|
wandb.log(to_report, commit=commit)
|