"""Predict the energy and the quality based on a neuronal network."""
from mendevi.models import Model
import torch
import tqdm
class _TorchNN(torch.nn.Module):
"""Simple perceptron nn."""
def __init__(self, in_ndim: int, out_ndim: int, total: int) -> None:
# verifications
assert isinstance(in_ndim, int), in_ndim.__class__.__name__
assert in_ndim >= 1, in_ndim
assert isinstance(out_ndim, int), out_ndim.__class__.__name__
assert out_ndim >= 1, out_ndim
assert isinstance(total, int), total.__class__.__name__
assert total >= in_ndim + 8 + out_ndim, total
super().__init__()
torch.manual_seed(0) # for reproductibility
# solve total = in_ndim*fact + fact * 2*fact + 2*fact * fact + fact*out_ndim
delta = (in_ndim + out_ndim)**2 + 16*total
fact = (delta**.5 - in_ndim - out_ndim) / 8
fact = round(fact)
# 4 full connected layers (better without dropout)
self.linear = torch.nn.Sequential(
torch.nn.Linear(in_ndim, fact),
torch.nn.ELU(),
torch.nn.Linear(fact, 2*fact),
torch.nn.BatchNorm1d(2*fact), # just this batch norm, no more
torch.nn.ELU(),
torch.nn.Linear(2*fact, fact),
torch.nn.ELU(),
torch.nn.Linear(fact, out_ndim),
)
def forward(self, values: torch.Tensor) -> torch.Tensor:
"""Inference of the vectorised model.
Parameters
----------
values : torch.Tensor
The concatenated values of shape (batch, in_ndim)
"""
assert isinstance(values, torch.Tensor), values.__class__.__name__
assert values.ndim == 2, values.shape
return self.linear(values)
[docs]
class DNN(Model):
"""Deep neuronal network predictive model."""
def _fit_vect_norm(self, values: dict[str, torch.Tensor]) -> object:
"""Perform the gradient descent on a just intanciated nn."""
# intitialisation
input_values = torch.cat([values[lbl] for lbl in self.input_labels_aggreg], dim=1)
reference = torch.cat([values[lbl] for lbl in self.output_labels], dim=1)
total = max(
input_values.shape[1] + 8 + reference.shape[1], # minimum structural
input_values.shape[0] // 5, # empirical rule
)
model = _TorchNN(input_values.shape[1], reference.shape[1], total=total)
# train
model.train()
optimizer = torch.optim.Adam(model.parameters(), lr=2e-4)
iterations = 10_000
with tqdm.tqdm(
desc="Gradient decay",
dynamic_ncols=True,
leave=False,
smoothing=1e-6,
total=iterations,
) as progress_bar:
for _ in range(iterations):
optimizer.zero_grad()
predict = model(input_values)
# handle nan values
nanref = reference.clone()
nanmask = nanref.isnan()
nanref[nanmask] = predict[nanmask]
loss = ((predict - nanref)**2).mean()
loss.backward()
optimizer.step()
progress_bar.update(1)
progress_bar.set_postfix_str(f"loss: {loss:9.3e}")
model.eval()
# serialisation
return (model, {lbl: values[lbl].shape[1] for lbl in self.output_labels})
def _predict_vect_norm(
self, values: dict[str, torch.Tensor], parameters: object,
) -> dict[str, torch.Tensor]:
"""Inference part of the model."""
# preparation, vectorisation
input_values = torch.cat([values[lbl] for lbl in self.input_labels_aggreg], dim=1)
model, sizes = parameters
# inference
prediction = model(input_values)
# parse result
output = {}
index = 0
for lbl in self.output_labels:
output[lbl] = prediction[:, index:index+sizes[lbl]]
index += sizes[lbl]
return output
# from mendevi.models.dnn import DNN_Quality
# DNN_Quality().fit("/data/dataset/merge.db", table="t_met_metric", select="lpips is not None and uniform(vid_hash) > 0.2").validate("/data/dataset/merge.db", table="t_met_metric", select="lpips is not None and uniform(vid_hash) <= 0.2")
[docs]
class DNN_Quality(DNN):
"""Deep neuronal network for the video quality prediction.
Examples
--------
>>> from mendevi.models.dnn import DNN_Quality
>>> model = DNN_Quality(output_labels=["psnr", "log_rev_ssim"])
>>> model.fit("svtav1_vs_rav1e_vs_aom.db", select="uniform(vid_hash) >= 0.2")
>>> model.validate("svtav1_vs_rav1e_vs_aom.db", select="uniform(vid_hash) < 0.2")
>>>
"""
def __init__(
self,
title: str = "Quality estimation based on a neuronal network",
**kwargs: dict,
) -> None:
"""Initialise the model."""
kwargs = kwargs.copy()
kwargs["input_labels"] = kwargs.get(
"input_labels", [
"encoder", "effort", "mode", "quality",
"nb_frames", "height", "width", # "range", "primaries", "eotf",
"spatial_dct", "temporal_dct", "rms_sobel", "rms_time_diff",
],
)
kwargs["output_labels"] = kwargs.get(
"output_labels",
["lpips", "psnr", "ssim", "vif", "vmaf", "log_size_per_frame"],
)
kwargs["aggregation"] = kwargs.get("aggregation", [])
super().__init__(title=title, **kwargs)
[docs]
class DNN_EncodeEnergy(DNN):
"""Deep neuronal network for the video encoding energy prediction.
Examples
--------
>>> from mendevi.models.dnn import DNN_EncodeEnergy
>>> model = DNN_EncodeEnergy().fit("svtav1_vs_rav1e_vs_aom.db", table="t_enc_encode")
>>> model.validate("svtav1_vs_rav1e_vs_aom.db")
>>>
"""
def __init__(
self,
title: str = "Encoding energy estimation based on a neuronal network",
**kwargs: dict,
) -> None:
"""Initialise the model."""
kwargs = kwargs.copy()
kwargs["input_labels"] = kwargs.get(
"input_labels", [
"encoder", "effort", "mode", "quality",
"hostname", # "temperature",
"nb_frames", "height", "width", # "range", "primaries", "eotf",
"spatial_dct", "temporal_dct", "rms_sobel", "rms_time_diff",
],
)
kwargs["output_labels"] = kwargs.get("output_labels", ["log_act_duration_per_frame", "power"])
kwargs["aggregation"] = kwargs.get("aggregation", ["hostname"])
super().__init__(title=title, **kwargs)
[docs]
class DNN_DecodeEnergy(DNN):
"""Deep neuronal network for the video encoding energy prediction.
Examples
--------
>>> from mendevi.models.dnn import DNN_DecodeEnergy
>>> model = DNN_DecodeEnergy().fit("svtav1_vs_rav1e_vs_aom.db", table="t_dec_decode")
>>> model.validate("svtav1_vs_rav1e_vs_aom.db")
>>>
"""
def __init__(
self,
title: str = "Decoding energy estimation based on a neuronal network",
**kwargs: dict,
) -> None:
"""Initialise the model."""
kwargs = kwargs.copy()
kwargs["input_labels"] = kwargs.get(
"input_labels", [
"decoder",
"hostname", # "temperature",
"nb_frames", "height", "width", # "range", "primaries", "eotf",
"spatial_dct", "temporal_dct", "rms_sobel", "rms_time_diff",
],
)
kwargs["output_labels"] = kwargs.get("output_labels", ["log_act_duration_per_frame", "power"])
kwargs["aggregation"] = kwargs.get("aggregation", ["hostname"])
super().__init__(title=title, **kwargs)