fei-nn-1.0.0: src/MXNet/NN/EvalMetric.hs
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeOperators #-}
module MXNet.NN.EvalMetric where
import Formatting (fixed, int, sformat, stext, (%))
import RIO
import qualified RIO.HashMap as M
import qualified RIO.HashMap.Partial as M ((!))
import qualified RIO.NonEmpty as RNE
import qualified RIO.Text as T
import qualified RIO.Vector.Storable as SV
import qualified RIO.Vector.Storable.Partial as SV (head)
import MXNet.Base
import MXNet.Base.Operators.Tensor (_norm)
import MXNet.NN.Layer
import MXNet.NN.Types
-- | Abstract Evaluation type class
class EvalMetricMethod metric where
data MetricData metric a
newMetric :: (MonadIO m, FloatDType a, HasCallStack)
=> Text -- phase name
-> metric a -- tag
-> m (MetricData metric a)
metricUpdate :: (MonadIO m, FloatDType a, HasCallStack)
=> MetricData metric a -- evaluation metric
-> M.HashMap Text (NDArray a) -- network bindings
-> [NDArray a] -- output of the network
-> m (M.HashMap Text Double)
metricName :: MetricData metric a -> Text
metricValue :: (MonadIO m, FloatDType a, HasCallStack) => MetricData metric a -> m Double
metricFormat :: (MonadIO m, FloatDType a, HasCallStack) => MetricData metric a -> m Text
metricFormat m = do
name <- pure (metricName m)
value <- metricValue m
return $ sformat ("<" % stext % ": " % fixed 4 % ">") name value
-- | Basic evaluation - accuracy
data AccuracyPredType = PredByThreshold Float
| PredByArgmax
| PredByArgmaxAt Int
data Accuracy a = Accuracy
{ _mtr_acc_name :: Maybe Text
, _mtr_acc_type :: AccuracyPredType
, _mtr_acc_min_value :: Float -- | to filter values less than min
, _mtr_acc_get_prob :: M.HashMap Text (NDArray a) -> [NDArray a] -> NDArray a
, _mtr_acc_get_gt :: M.HashMap Text (NDArray a) -> [NDArray a] -> NDArray a
}
instance EvalMetricMethod Accuracy where
data MetricData Accuracy a = AccuracyPriv (Accuracy a) Text (IORef Int) (IORef Int)
newMetric phase conf = do
a <- liftIO $ newIORef 0
b <- liftIO $ newIORef 0
return $ AccuracyPriv conf phase a b
metricUpdate mtr@(AccuracyPriv Accuracy{..} phase cntRef sumRef) bindings outputs = liftIO $ do
out <- toCPU $ _mtr_acc_get_prob bindings outputs
lbl <- toCPU $ _mtr_acc_get_gt bindings outputs
out <- case _mtr_acc_type of
PredByThreshold thr -> geqScalar thr out
PredByArgmax -> argmax out (Just (-1))False
PredByArgmaxAt axis -> argmax out (Just axis) False
valid <- geqScalar _mtr_acc_min_value lbl
correct <- and_ valid =<< eq_ out lbl
num_correct <- SV.head <$> (toVector =<< sum_ correct Nothing False)
num_valid <- SV.head <$> (toVector =<< sum_ valid Nothing False)
modifyIORef sumRef (+ floor num_correct)
modifyIORef cntRef (+ floor num_valid)
value <- metricValue mtr
return $ M.singleton (metricName mtr) value
metricName (AccuracyPriv Accuracy{..} phase _ _) =
let name = fromMaybe "acc" _mtr_acc_name
in sformat (stext % "_" % stext) phase name
metricValue (AccuracyPriv _ _ cntRef sumRef) = liftIO $ do
s <- liftIO $ readIORef sumRef
n <- liftIO $ readIORef cntRef
return (100 * fromIntegral s / fromIntegral n)
-- | Basic evaluation - vector norm
data Norm a = Norm
{ _mtr_norm_name :: Maybe Text
, _mtr_norm_ord :: Int
, _mtr_norm_get_array :: M.HashMap Text (NDArray a) -> [NDArray a] -> NDArray a
}
instance EvalMetricMethod Norm where
data MetricData Norm a = NormPriv (Norm a) Text (IORef Int) (IORef Double)
newMetric phase conf = do
a <- liftIO $ newIORef 0
b <- liftIO $ newIORef 0
return $ NormPriv conf phase a b
metricUpdate mtr@(NormPriv Norm{..} phase cntRef sumRef) bindings preds = liftIO $ do
array <- toCPU $ _mtr_norm_get_array bindings preds
norm <- prim _norm (#data := array .& #ord := _mtr_norm_ord .& Nil)
norm <- SV.head <$> toVector norm
batch_size :| _ <- ndshape array
modifyIORef' sumRef (+ realToFrac norm)
modifyIORef' cntRef (+ batch_size)
value <- metricValue mtr
return $ M.singleton (metricName mtr) value
metricName (NormPriv Norm{..} phase _ _) =
let lk = sformat ("_L" % int) _mtr_norm_ord
name = fromMaybe lk _mtr_norm_name
in sformat (stext % "_" % stext) phase name
metricValue (NormPriv _ _ cntRef sumRef) = liftIO $ do
s <- readIORef sumRef
n <- readIORef cntRef
return $ realToFrac s / fromIntegral n
-- | Basic evaluation - cross entropy
data CrossEntropy a = CrossEntropy
{ _mtr_ce_name :: Maybe Text
, _mtr_ce_gt_clsid :: Bool
, _mtr_ce_get_prob :: M.HashMap Text (NDArray a) -> [NDArray a] -> NDArray a
, _mtr_ce_get_gt :: M.HashMap Text (NDArray a) -> [NDArray a] -> NDArray a
}
instance EvalMetricMethod CrossEntropy where
data MetricData CrossEntropy a = CrossEntropyPriv (CrossEntropy a) Text (IORef Int) (IORef Double)
newMetric phase conf = do
a <- liftIO $ newIORef 0
b <- liftIO $ newIORef 0
return $ CrossEntropyPriv conf phase a b
metricUpdate mtr@(CrossEntropyPriv CrossEntropy{..} phase cntRef sumRef) bindings preds = liftIO $ do
prob <- toCPU $ _mtr_ce_get_prob bindings preds
gt <- toCPU $ _mtr_ce_get_gt bindings preds
(loss, num_valid) <-
if _mtr_ce_gt_clsid
then do
-- when the gt labels are class id,
-- prob: (a, ..., b, num_classes)
-- gt: (a,..,b)
-- dim(prob) = dim(gt) + 1
-- The last dimension serves as the prob dist.
-- We pickup the prob at the label specified class.
cls_prob <- log2_ =<< addScalar 1e-5 =<< pickI gt prob
weights <- geqScalar 0 gt
cls_prob <- mul_ cls_prob weights
nloss <- toVector =<< sum_ cls_prob Nothing False
num_valid <- toVector =<< sum_ weights Nothing False
return (negate (SV.head nloss), SV.head num_valid)
else do
-- when the gt are onehot vector
-- prob: (a, .., b, num_classes)
-- gt: (a, .., b, num_classes)
-- dim(prob) == dim(gt)
term1 <- mul_ gt =<< log2_ =<< addScalar 1e-5 prob
a <- log2_ =<< addScalar 1e-5 =<< rsubScalar 1 prob
term2 <- mul_ a =<< rsubScalar 1 gt
weights <- geqScalar 0 gt
nloss <- mul_ weights =<< add_ term1 term2
nloss <- toVector =<< sum_ nloss Nothing False
num_valid <- toVector =<< sum_ weights Nothing False
return (negate (SV.head nloss), SV.head num_valid)
modifyIORef' sumRef (+ realToFrac loss)
modifyIORef' cntRef (+ floor num_valid)
value <- metricValue mtr
return $ M.singleton (metricName mtr) value
metricName (CrossEntropyPriv CrossEntropy{..} phase _ _) =
let name = fromMaybe "ce" _mtr_ce_name
in sformat (stext % "_" % stext) phase name
metricValue (CrossEntropyPriv _ _ cntRef sumRef) = liftIO $ do
s <- readIORef sumRef
n <- readIORef cntRef
return $ realToFrac s / fromIntegral n
data ListOfMetric ms a where
MNil :: ListOfMetric '[] a
(:*) :: (EvalMetricMethod m) => m a -> ListOfMetric ms a -> ListOfMetric (m ': ms) a
instance EvalMetricMethod (ListOfMetric '[]) where
data MetricData (ListOfMetric '[]) a = MNilData
newMetric _ _ = return MNilData
metricName _ = error "Empty metric"
metricValue _ = error "Empty metric"
metricUpdate _ _ _ = return M.empty
metricFormat _ = return ""
instance (EvalMetricMethod m, EvalMetricMethod (ListOfMetric ms)) => EvalMetricMethod (ListOfMetric (m ': ms)) where
data MetricData (ListOfMetric (m ': ms)) a = MCompositeData (MetricData m a) (MetricData (ListOfMetric ms) a)
newMetric phase (a :* as) = MCompositeData <$> (newMetric phase a) <*> (newMetric phase as)
metricUpdate (MCompositeData a as) bindings output = do
m1 <- metricUpdate a bindings output
m2 <- metricUpdate as bindings output
return $ M.union m1 m2
metricName _ = error "List of metrics"
metricValue _ = error "List of metrics"
metricFormat (MCompositeData a as) = do
s1 <- metricFormat a
s2 <- metricFormat as
return $ T.concat [s1, " ", s2]
infixr 9 :*
class MetricsToList ms where
metricsToList :: (MonadIO m, FloatDType a) => MetricData (ListOfMetric ms) a -> m [(Text, Double)]
instance MetricsToList '[] where
metricsToList MNilData = return []
instance (EvalMetricMethod m, MetricsToList n) => MetricsToList (m ': n) where
metricsToList (MCompositeData a b) = do
n <- pure $ metricName a
v <- metricValue a
w <- metricsToList b
return $ (n, v) : w