46 lines
1.6 KiB
Python
46 lines
1.6 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
import torch
|
||
|
|
|
||
|
|
|
||
|
|
def count_parameters(model: torch.nn.Module) -> int:
|
||
|
|
return sum(parameter.numel() for parameter in model.parameters() if parameter.requires_grad)
|
||
|
|
|
||
|
|
|
||
|
|
def per_channel_mse(
|
||
|
|
squared_error_sum: torch.Tensor,
|
||
|
|
count: int,
|
||
|
|
target_names: tuple[str, ...],
|
||
|
|
) -> dict[str, float]:
|
||
|
|
if count <= 0:
|
||
|
|
raise ValueError("Metric count must be positive")
|
||
|
|
values = (squared_error_sum / count).detach().cpu().tolist()
|
||
|
|
return {name: float(value) for name, value in zip(target_names, values, strict=True)}
|
||
|
|
|
||
|
|
|
||
|
|
def overall_mse(squared_error_sum: torch.Tensor, count: int, target_dim: int) -> float:
|
||
|
|
if count <= 0 or target_dim <= 0:
|
||
|
|
raise ValueError("Metric count and target dimension must be positive")
|
||
|
|
return float((squared_error_sum.sum() / (count * target_dim)).detach().cpu().item())
|
||
|
|
|
||
|
|
|
||
|
|
def device_metrics(device: torch.device) -> dict[str, Any]:
|
||
|
|
if device.type != "cuda":
|
||
|
|
return {
|
||
|
|
"device": str(device),
|
||
|
|
"gpu_name": None,
|
||
|
|
"gpu_memory_total_mb": None,
|
||
|
|
"gpu_memory_peak_allocated_mb": None,
|
||
|
|
}
|
||
|
|
|
||
|
|
index = device.index if device.index is not None else torch.cuda.current_device()
|
||
|
|
properties = torch.cuda.get_device_properties(index)
|
||
|
|
return {
|
||
|
|
"device": f"cuda:{index}",
|
||
|
|
"gpu_name": torch.cuda.get_device_name(index),
|
||
|
|
"gpu_memory_total_mb": int(properties.total_memory // (1024 * 1024)),
|
||
|
|
"gpu_memory_peak_allocated_mb": int(torch.cuda.max_memory_allocated(index) // (1024 * 1024)),
|
||
|
|
}
|