asphodel-thuang's picture
Upload 59 files
ae522ad verified
import numpy as np
import mlflow
from mwm.components.metrics import *
class MetricsLogger2Channel:
def __init__(self):
self.epoch_loss = 0.0
self.epoch_loss_val = 0.0
self.epoch_iou_ch0 = 0.0
self.epoch_iou_ch1 = 0.0
def update_sum(self, loss, preds, masks):
self.epoch_loss += loss.item()
self.epoch_iou_ch0 += np.sum(iou_list(preds[:,0,:,:], masks[:,0,:,:]))
self.epoch_iou_ch1 += np.sum(iou_list(preds[:,1,:,:], masks[:,1,:,:]))
def update_mean(self, num_batch, num_samples):
self.epoch_loss /= num_batch
self.epoch_iou_ch0 /= num_samples
self.epoch_iou_ch1 /= num_samples
def update_sum_val(self, loss):
self.epoch_loss_val += loss.item()
def update_mean_val(self, num_batch):
self.epoch_loss_val /= num_batch
def log_metrics_mlflow(self, epoch):
mlflow.log_metric("loss_train", self.epoch_loss, step=epoch)
mlflow.log_metric("loss_val", self.epoch_loss_val, step=epoch)
mlflow.log_metric("iou_ch0", self.epoch_iou_ch0, step=epoch)
mlflow.log_metric("iou_ch1", self.epoch_iou_ch1, step=epoch)
print(f"Epoch {epoch}, \
Loss: {self.epoch_loss:.4f}, \
Val Loss: {self.epoch_loss_val:.4f}, \
Avg IoU Ch0: {self.epoch_iou_ch0:.4f}, \
Avg IoU Ch1: {self.epoch_iou_ch1:.4f}")
self.reset()
def reset(self):
self.epoch_loss = 0.0
self.epoch_loss_val = 0.0
self.epoch_iou_ch0 = 0.0
self.epoch_iou_ch1 = 0.0