def print_any_prediction(criterion, model, nested_loaders, runner):
loader = nested_loaders['valid']
X, y = next(iter(loader))
X = any2device(X, next(model.parameters()).device)
y = any2device(y, next(model.parameters()).device)
model = model.eval()
pred = model(X)
res = criterion(pred, y)
print(res.item())
print(pred, y)
def load_checkpoint(self,
*,
filepath: str = None,
db_server: DBSpec = None) -> bool:
if filepath is not None:
checkpoint = utils.load_checkpoint(filepath)
elif db_server is not None:
current_epoch = db_server.epoch
checkpoint = db_server.get_checkpoint()
if not db_server.training_enabled \
and db_server.epoch == current_epoch:
return False
while checkpoint is None or db_server.epoch <= current_epoch:
time.sleep(3.0)
checkpoint = db_server.get_checkpoint()
if not db_server.training_enabled \
and db_server.epoch == current_epoch:
return False
else:
return False
self.checkpoint = checkpoint
weights = self.checkpoint[f"{self._weights_sync_mode}_state_dict"]
weights = {
k: utils.any2device(v, device=self._device)
for k, v in weights.items()
}
self.agent.load_state_dict(weights)
self.agent.to(self._device)
self.agent.eval()
return True
def compute_predictions(model,
dataset,
batch_size=1,
workers=0) -> pd.DataFrame:
df = defaultdict(list)
for batch in tqdm(
DataLoader(dataset,
batch_size=batch_size,
num_workers=workers,
shuffle=False,
drop_last=False,
pin_memory=True)):
batch = any2device(batch, device="cuda")
image_ids = batch[INPUT_IMAGE_ID_KEY]
df[INPUT_IMAGE_ID_KEY].extend(image_ids)
outputs = model(**batch)
if OUTPUT_PRED_MODIFICATION_FLAG in outputs:
df[OUTPUT_PRED_MODIFICATION_FLAG].extend(
to_numpy(outputs[OUTPUT_PRED_MODIFICATION_FLAG]).flatten())
if OUTPUT_PRED_MODIFICATION_TYPE in outputs:
df[OUTPUT_PRED_MODIFICATION_TYPE].extend(
to_numpy(outputs[OUTPUT_PRED_MODIFICATION_TYPE]).tolist())
df = pd.DataFrame.from_dict(df)
return df
def compute_trn_predictions(model,
dataset,
fp16=False,
batch_size=1,
workers=0) -> pd.DataFrame:
df = defaultdict(list)
for batch in tqdm(
DataLoader(dataset,
batch_size=batch_size,
num_workers=workers,
shuffle=False,
drop_last=False,
pin_memory=True)):
batch = any2device(batch, device="cuda")
if fp16 and INPUT_FEATURES_JPEG_FLOAT in batch:
batch[INPUT_FEATURES_JPEG_FLOAT] = batch[
INPUT_FEATURES_JPEG_FLOAT].half()
if INPUT_TRUE_MODIFICATION_FLAG in batch:
y_trues = to_numpy(batch[INPUT_TRUE_MODIFICATION_FLAG]).flatten()
df[INPUT_TRUE_MODIFICATION_FLAG].extend(y_trues)
if INPUT_TRUE_MODIFICATION_TYPE in batch:
y_labels = to_numpy(batch[INPUT_TRUE_MODIFICATION_TYPE]).flatten()
df[INPUT_TRUE_MODIFICATION_TYPE].extend(y_labels)
image_ids = batch[INPUT_IMAGE_ID_KEY]
df[INPUT_IMAGE_ID_KEY].extend(image_ids)
outputs = model(**batch)
if OUTPUT_PRED_MODIFICATION_FLAG in outputs:
df[OUTPUT_PRED_MODIFICATION_FLAG].extend(
to_numpy(outputs[OUTPUT_PRED_MODIFICATION_FLAG]).flatten())
if OUTPUT_PRED_MODIFICATION_TYPE in outputs:
df[OUTPUT_PRED_MODIFICATION_TYPE].extend(
outputs[OUTPUT_PRED_MODIFICATION_TYPE].tolist())
if OUTPUT_PRED_EMBEDDING in outputs:
df[OUTPUT_PRED_EMBEDDING].extend(
outputs[OUTPUT_PRED_EMBEDDING].tolist())
# Save also TTA predictions for future use
if OUTPUT_PRED_MODIFICATION_FLAG + "_tta" in outputs:
df[OUTPUT_PRED_MODIFICATION_FLAG + "_tta"].extend(
to_numpy(outputs[OUTPUT_PRED_MODIFICATION_FLAG +
"_tta"]).tolist())
if OUTPUT_PRED_MODIFICATION_TYPE + "_tta" in outputs:
df[OUTPUT_PRED_MODIFICATION_TYPE + "_tta"].extend(
to_numpy(outputs[OUTPUT_PRED_MODIFICATION_TYPE +
"_tta"]).tolist())
df = pd.DataFrame.from_dict(df)
return df
def main(args, unknown_args):
args, config = parse_args_uargs(args, unknown_args)
set_global_seed(args.seed)
prepare_cudnn(args.deterministic, args.benchmark)
if args.logdir is not None:
os.makedirs(args.logdir, exist_ok=True)
dump_environment(config, args.logdir, args.configs)
if args.expdir is not None:
module = import_module(expdir=args.expdir) # noqa: F841
if args.logdir is not None:
dump_code(args.expdir, args.logdir)
env = ENVIRONMENTS.get_from_params(**config["environment"])
algorithm_name = config["algorithm"].pop("algorithm")
if algorithm_name in OFFPOLICY_ALGORITHMS_NAMES:
ALGORITHMS = OFFPOLICY_ALGORITHMS
trainer_fn = OffpolicyTrainer
sync_epoch = False
elif algorithm_name in ONPOLICY_ALGORITHMS_NAMES:
ALGORITHMS = ONPOLICY_ALGORITHMS
trainer_fn = OnpolicyTrainer
sync_epoch = True
else:
# @TODO: add registry for algorithms, trainers, samplers
raise NotImplementedError()
db_server = DATABASES.get_from_params(
**config.get("db", {}), sync_epoch=sync_epoch
)
algorithm_fn = ALGORITHMS.get(algorithm_name)
algorithm = algorithm_fn.prepare_for_trainer(env_spec=env, config=config)
if args.resume is not None:
checkpoint = utils.load_checkpoint(filepath=args.resume)
checkpoint = utils.any2device(checkpoint, utils.get_device())
algorithm.unpack_checkpoint(
checkpoint=checkpoint,
with_optimizer=False
)
monitoring_params = config.get("monitoring_params", None)
trainer = trainer_fn(
algorithm=algorithm,
env_spec=env,
db_server=db_server,
logdir=args.logdir,
monitoring_params=monitoring_params,
**config["trainer"],
)
trainer.run()
def get_criterion(self):
"""
Fetches the criterion. (Only one loss.)
"""
loss_name = self.criterion_params["loss"]
loss_kwargs = self.criterion_params[loss_name]
if "weight" in list(loss_kwargs.keys()):
if isinstance(loss_kwargs["weight"], list):
weight_tensor = torch.tensor(loss_kwargs["weight"])
weight_tensor = any2device(weight_tensor, get_device())
print(f"Converted the `weight` argument in {loss_name}",
f" to a {weight_tensor.type()}...")
loss_kwargs["weight"] = weight_tensor
loss_cls = globals()[loss_name]
loss = loss_cls(**loss_kwargs)
print(f"Criterion: {loss}")
return loss
def validate_model(runner, loader, pruning_fn, num_sessions):
accuracy_scores = []
pruned_weights = []
c_p = 1
for pruning_idx in range(num_sessions):
correct = 0
len_dataset = 0
for batch in loader:
outp = runner.predict_batch(utils.any2device(batch, "cuda"))
c_correct = torch.sum(
outp["logits"].argmax(-1).detach().cpu() == batch["targets"]
).item()
correct += c_correct
len_dataset += batch["features"].size(0)
pruned_weights.append(c_p)
c_p *= 0.9
accuracy_scores.append(correct / len_dataset)
pruning_fn(runner.model)
return accuracy_scores, pruned_weights
def bn_update(loader: DataLoader, model: nn.Module):
"""
BatchNorm buffers update (if any).
Performs 1 epochs to estimate buffers average using train dataset.
:param loader: train dataset loader for buffers average estimation.
:param model: model being update
:return: None
"""
if not check_bn(model):
return
assert loader.drop_last
model.train()
model.apply(reset_bn)
for batch in tqdm(loader, desc="AdaBN"):
batch = any2device(batch, device="cuda")
model(**batch)
model.apply(fix_bn)
def load_checkpoint(
self, *, filepath: str = None, db_server: DBSpec = None
):
if filepath is not None:
checkpoint = utils.load_checkpoint(filepath)
elif db_server is not None:
checkpoint = db_server.load_checkpoint()
while checkpoint is None:
time.sleep(3.0)
checkpoint = db_server.load_checkpoint()
else:
raise NotImplementedError
self.checkpoint = checkpoint
weights = self.checkpoint[f"{self._weights_sync_mode}_state_dict"]
weights = {
k: utils.any2device(v, device=self._device)
for k, v in weights.items()
}
self.agent.load_state_dict(weights)
self.agent.to(self._device)
self.agent.eval()
def train_fn(epoch, train_dataloader, optimizer, criterion, scheduler,
device):
model.train()
for batch_idx, batch_data in enumerate(train_dataloader):
optimizer.zero_grad()
batch_data = any2device(batch_data, device)
outputs = model(**batch_data)
y_pred = outputs[OUTPUT_PRED_MODIFICATION_TYPE]
y_true = batch_data[INPUT_TRUE_MODIFICATION_TYPE]
loss = criterion(y_pred, y_true)
if batch_idx % 100:
xm.master_print(f"Batch: {batch_idx}, loss: {loss.item()}")
loss.backward()
xm.optimizer_step(optimizer)
if scheduler is not None:
scheduler.step()
def valid_fn(epoch, valid_dataloader, criterion, device):
model.eval()
pred_scores = []
true_scores = []
for batch_idx, batch_data in enumerate(valid_dataloader):
batch_data = any2device(batch_data, device)
outputs = model(**batch_data)
y_pred = outputs[OUTPUT_PRED_MODIFICATION_TYPE]
y_true = batch_data[INPUT_TRUE_MODIFICATION_TYPE]
loss = criterion(y_pred, y_true)
pred_scores.extend(to_numpy(parse_classifier_probas(y_pred)))
true_scores.extend(to_numpy(y_true))
xm.master_print(f"Batch: {batch_idx}, loss: {loss.item()}")
val_wauc = alaska_weighted_auc(xla_all_gather(true_scores, device),
xla_all_gather(pred_scores, device))
xm.master_print(f"Valid epoch: {epoch}, wAUC: {val_wauc}")
return val_wauc
def _batch2device(self, batch: Mapping[str, Any], device: Device):
output = utils.any2device(batch, device)
return output
def trace_model_from_runner(
runner: IRunner,
checkpoint_name: str = None,
method_name: str = "forward",
mode: str = "eval",
requires_grad: bool = False,
opt_level: str = None,
device: Device = "cpu",
) -> ScriptModule:
"""
Traces model using created experiment and runner.
Args:
runner (Runner): Current runner.
checkpoint_name (str): Name of model checkpoint to use, if None
traces current model from runner
method_name (str): Model's method name that will be
used as entrypoint during tracing
mode (str): Mode for model to trace (``train`` or ``eval``)
requires_grad (bool): Flag to use grads
opt_level (str): AMP FP16 init level
device (str): Torch device
Returns:
(ScriptModule): Traced model
"""
logdir = runner.logdir
model = get_nn_from_ddp_module(runner.model)
if checkpoint_name is not None:
dumped_checkpoint = pack_checkpoint(model=model)
checkpoint_path = logdir / "checkpoints" / f"{checkpoint_name}.pth"
checkpoint = load_checkpoint(filepath=checkpoint_path)
unpack_checkpoint(checkpoint=checkpoint, model=model)
# getting input names of args for method since we don't have Runner
# and we don't know input_key to preprocess batch for method call
fn = getattr(model, method_name)
method_argnames = _get_input_argnames(fn=fn, exclude=["self"])
batch = {}
for name in method_argnames:
# TODO: We don't know input_keys without runner
assert name in runner.input, (
"Input batch should contain the same keys as input argument "
"names of `forward` function to be traced correctly")
batch[name] = runner.input[name]
batch = any2device(batch, device)
# Dumping previous runner of the model, we will need it to restore
_device, _is_training, _requires_grad = (
runner.device,
model.training,
get_requires_grad(model),
)
model.to(device)
# Function to run prediction on batch
def predict_fn(model: Model, inputs, **kwargs):
return model(**inputs, **kwargs)
traced_model = trace_model(
model=model,
predict_fn=predict_fn,
batch=batch,
method_name=method_name,
mode=mode,
requires_grad=requires_grad,
opt_level=opt_level,
device=device,
)
if checkpoint_name is not None:
unpack_checkpoint(checkpoint=dumped_checkpoint, model=model)
# Restore previous runner of the model
getattr(model, "train" if _is_training else "eval")()
set_requires_grad(model, _requires_grad)
model.to(_device)
return traced_model
def batch_to_model_device(self, batch) -> Mapping[str, torch.Tensor]:
return any2device(batch, next(self.model.parameters()).device)
def batch_to_device(self, batch, device) -> Mapping[str, torch.Tensor]:
return any2device(batch, device)
def main(args, _=None):
"""Run the ``catalyst-data text2embeddings`` script."""
batch_size = args.batch_size
num_workers = args.num_workers
max_length = args.max_length
pooling_groups = args.pooling.split(",")
bert_level = args.bert_level
if bert_level is not None:
assert (args.output_hidden_states
), "You need hidden states output for level specification"
utils.set_global_seed(args.seed)
utils.prepare_cudnn(args.deterministic, args.benchmark)
if getattr(args, "in_huggingface", False):
model_config = BertConfig.from_pretrained(args.in_huggingface)
model_config.output_hidden_states = args.output_hidden_states
model = BertModel.from_pretrained(args.in_huggingface,
config=model_config)
tokenizer = BertTokenizer.from_pretrained(args.in_huggingface)
else:
model_config = BertConfig.from_pretrained(args.in_config)
model_config.output_hidden_states = args.output_hidden_states
model = BertModel(config=model_config)
tokenizer = BertTokenizer.from_pretrained(args.in_vocab)
if getattr(args, "in_model", None) is not None:
checkpoint = utils.load_checkpoint(args.in_model)
checkpoint = {"model_state_dict": checkpoint}
utils.unpack_checkpoint(checkpoint=checkpoint, model=model)
model = model.eval()
model, _, _, _, device = utils.process_components(model=model)
df = pd.read_csv(args.in_csv)
df = df.dropna(subset=[args.txt_col])
df.to_csv(f"{args.out_prefix}.df.csv", index=False)
df = df.reset_index().drop("index", axis=1)
df = list(df.to_dict("index").values())
num_samples = len(df)
open_fn = LambdaReader(
input_key=args.txt_col,
output_key=None,
lambda_fn=partial(
tokenize_text,
strip=args.strip,
lowercase=args.lowercase,
remove_punctuation=args.remove_punctuation,
),
tokenizer=tokenizer,
max_length=max_length,
)
dataloader = utils.get_loader(
df,
open_fn,
batch_size=batch_size,
num_workers=num_workers,
)
features = {}
dataloader = tqdm(dataloader) if args.verbose else dataloader
with torch.no_grad():
for idx, batch_input in enumerate(dataloader):
batch_input = utils.any2device(batch_input, device)
batch_output = model(**batch_input)
mask = (batch_input["attention_mask"].unsqueeze(-1)
if args.mask_for_max_length else None)
if utils.check_ddp_wrapped(model):
# using several gpu
hidden_size = model.module.config.hidden_size
hidden_states = model.module.config.output_hidden_states
else:
# using cpu or one gpu
hidden_size = model.config.hidden_size
hidden_states = model.config.output_hidden_states
batch_features = process_bert_output(
bert_output=batch_output,
hidden_size=hidden_size,
output_hidden_states=hidden_states,
pooling_groups=pooling_groups,
mask=mask,
)
# create storage based on network output
if idx == 0:
for layer_name, layer_value in batch_features.items():
if bert_level is not None and bert_level != layer_name:
continue
layer_name = (layer_name if isinstance(layer_name, str)
else f"{layer_name:02d}")
_, embedding_size = layer_value.shape
features[layer_name] = np.memmap(
f"{args.out_prefix}.{layer_name}.npy",
dtype=np.float32,
mode="w+",
shape=(num_samples, embedding_size),
)
indices = np.arange(idx * batch_size,
min((idx + 1) * batch_size, num_samples))
for layer_name2, layer_value2 in batch_features.items():
if bert_level is not None and bert_level != layer_name2:
continue
layer_name2 = (layer_name2 if isinstance(layer_name2, str) else
f"{layer_name2:02d}")
features[layer_name2][indices] = _detach(layer_value2)
if args.force_save:
for key, mmap in features.items():
mmap.flush()
np.save(f"{args.out_prefix}.{key}.force.npy",
mmap,
allow_pickle=False)
