Repository: deepjavalibrary/djl
Pull request #1613 (2022-04-29 05:10:24)

[api] Refactor metrics API
Change-Id: I271ec2b5b89034d837949eee4829016321550e34

## Description ##

Brief description of what this PR is about

- If this change is a backward incompatible change, why must this change be made?
- Interesting edge cases to note here


codecov-commenter: # [Codecov](https://codecov.io/gh/deepjavalibrary/djl/pull/1613?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) Report
> Merging [#1613](https://codecov.io/gh/deepjavalibrary/djl/pull/1613?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) (73edd2a) into [master](https://codecov.io/gh/deepjavalibrary/djl/commit/bb5073f7ffc33f1883621fc1783dbb01e70b418f?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) (bb5073f) will **decrease** coverage by `1.21%`.
> The diff coverage is `61.87%`.

```diff
@@             Coverage Diff              @@
##             master    #1613      +/-   ##
============================================
- Coverage     72.08%   70.87%   -1.22%     
- Complexity     5126     5420     +294     
============================================
  Files           473      506      +33     
  Lines         21970    23712    +1742     
  Branches       2351     2582     +231     
============================================
+ Hits          15838    16805     +967     
- Misses         4925     5621     +696     
- Partials       1207     1286      +79     
```


| [Impacted Files](https://codecov.io/gh/deepjavalibrary/djl/pull/1613?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | Coverage Δ | |
|---|---|---|
| [api/src/main/java/ai/djl/modality/cv/Image.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1613/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L0ltYWdlLmphdmE=) | `69.23% <ø> (-4.11%)` | :arrow_down: |
| [...i/djl/modality/cv/translator/BigGANTranslator.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1613/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvQmlnR0FOVHJhbnNsYXRvci5qYXZh) | `21.42% <ø> (-5.24%)` | :arrow_down: |
| [...odality/cv/translator/BigGANTranslatorFactory.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1613/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvQmlnR0FOVHJhbnNsYXRvckZhY3RvcnkuamF2YQ==) | `33.33% <0.00%> (+8.33%)` | :arrow_up: |
| [...nslator/InstanceSegmentationTranslatorFactory.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1613/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvSW5zdGFuY2VTZWdtZW50YXRpb25UcmFuc2xhdG9yRmFjdG9yeS5qYXZh) | `14.28% <0.00%> (-3.90%)` | :arrow_down: |
| [.../modality/cv/translator/YoloTranslatorFactory.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1613/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvWW9sb1RyYW5zbGF0b3JGYWN0b3J5LmphdmE=) | `8.33% <0.00%> (-1.67%)` | :arrow_down: |
| [...i/djl/modality/cv/translator/YoloV5Translator.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1613/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvWW9sb1Y1VHJhbnNsYXRvci5qYXZh) | `5.69% <0.00%> (ø)` | |
| [...odality/cv/translator/YoloV5TranslatorFactory.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1613/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvWW9sb1Y1VHJhbnNsYXRvckZhY3RvcnkuamF2YQ==) | `8.33% <0.00%> (-1.67%)` | :arrow_down: |
| [...pi/src/main/java/ai/djl/ndarray/BytesSupplier.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1613/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL25kYXJyYXkvQnl0ZXNTdXBwbGllci5qYXZh) | `54.54% <0.00%> (-12.13%)` | :arrow_down: |
| [...ain/java/ai/djl/ndarray/index/dim/NDIndexPick.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1613/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL25kYXJyYXkvaW5kZXgvZGltL05ESW5kZXhQaWNrLmphdmE=) | `100.00% <ø> (ø)` | |
| [api/src/main/java/ai/djl/nn/Blocks.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1613/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL25uL0Jsb2Nrcy5qYXZh) | `75.00% <0.00%> (-25.00%)` | :arrow_down: |
| ... and [239 more](https://codecov.io/gh/deepjavalibrary/djl/pull/1613/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/deepjavalibrary/djl/pull/1613?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/deepjavalibrary/djl/pull/1613?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary). Last update [8ff1b4f...73edd2a](https://codecov.io/gh/deepjavalibrary/djl/pull/1613?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).


Patch:
diff --git a/api/src/main/java/ai/djl/inference/Predictor.java b/api/src/main/java/ai/djl/inference/Predictor.java
index 7808b8cf2..1e234bf74 100644
--- a/api/src/main/java/ai/djl/inference/Predictor.java
+++ b/api/src/main/java/ai/djl/inference/Predictor.java
@@ -15,6 +15,7 @@ package ai.djl.inference;
 import ai.djl.Device;
 import ai.djl.Model;
 import ai.djl.metric.Metrics;
+import ai.djl.metric.Unit;
 import ai.djl.ndarray.LazyNDArray;
 import ai.djl.ndarray.NDArray;
 import ai.djl.ndarray.NDList;
@@ -228,9 +229,9 @@ public class Predictor<I, O> implements AutoCloseable {
         if (metrics != null) {
             waitToRead(list);
             long tmp = System.nanoTime();
-            long duration = tmp - timestamp;
+            long duration = (tmp - timestamp) / 1000;
             timestamp = tmp;
-            metrics.addMetric("Preprocess", duration, "nano");
+            metrics.addMetric("Preprocess", duration, Unit.MICROSECONDS);
         }
     }
 
@@ -238,19 +239,19 @@ public class Predictor<I, O> implements AutoCloseable {
         if (metrics != null) {
             waitToRead(list);
             long tmp = System.nanoTime();
-            long duration = tmp - timestamp;
+            long duration = (tmp - timestamp) / 1000;
             timestamp = tmp;
-            metrics.addMetric("Inference", duration, "nano");
+            metrics.addMetric("Inference", duration, Unit.MICROSECONDS);
         }
     }
 
     private void postProcessEnd(long begin) {
         if (metrics != null) {
             long tmp = System.nanoTime();
-            long duration = tmp - timestamp;
+            long duration = (tmp - timestamp) / 1000;
             timestamp = tmp;
-            metrics.addMetric("Postprocess", duration, "nano");
-            metrics.addMetric("Total", tmp - begin, "nano");
+            metrics.addMetric("Postprocess", duration, Unit.MICROSECONDS);
+            metrics.addMetric("Total", (tmp - begin) / 1000, Unit.MICROSECONDS);
         }
     }
 
diff --git a/api/src/main/java/ai/djl/metric/Dimension.java b/api/src/main/java/ai/djl/metric/Dimension.java
new file mode 100644
index 000000000..59028aeb7
--- /dev/null
+++ b/api/src/main/java/ai/djl/metric/Dimension.java
@@ -0,0 +1,57 @@
+/*
+ * Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.metric;
+
+import com.google.gson.annotations.SerializedName;
+
+/** A class represents a metric dimension. */
+public class Dimension {
+
+    @SerializedName("Name")
+    private String name;
+
+    @SerializedName("Value")
+    private String value;
+
+    /** Constructs a new {@code Dimension} instance. */
+    public Dimension() {}
+
+    /**
+     * Constructs a new {@code Dimension} instance.
+     *
+     * @param name the dimension name
+     * @param value the dimension value
+     */
+    public Dimension(String name, String value) {
+        this.name = name;
+        this.value = value;
+    }
+
+    /**
+     * Returns the dimension name.
+     *
+     * @return the dimension name
+     */
+    public String getName() {
+        return name;
+    }
+
+    /**
+     * Returns the dimension value.
+     *
+     * @return the dimension value
+     */
+    public String getValue() {
+        return value;
+    }
+}
diff --git a/api/src/main/java/ai/djl/metric/Metric.java b/api/src/main/java/ai/djl/metric/Metric.java
index 9a2cde566..1a8c44894 100644
--- a/api/src/main/java/ai/djl/metric/Metric.java
+++ b/api/src/main/java/ai/djl/metric/Metric.java
@@ -12,6 +12,14 @@
  */
 package ai.djl.metric;
 
+import com.google.gson.annotations.SerializedName;
+import java.net.InetAddress;
+import java.net.UnknownHostException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+
 /**
  * A class representing a single recorded {@code Metric} value.
  *
@@ -19,10 +27,29 @@ package ai.djl.metric;
  */
 public class Metric {
 
+    private static final Pattern PATTERN =
+            Pattern.compile(
+                    "\\s*([\\w\\s]+)\\.([\\w\\s]+):([0-9\\-,.e]+)\\|#([^|]*)\\|#hostname:([^,]+),([^,]+)(,.*)?");
+
+    private static final String LOCALHOST = getLocalHostName();
+
+    @SerializedName("MetricName")
     private String metricName;
-    private Number value;
+
+    @SerializedName("Value")
+    private String value;
+
+    @SerializedName("Unit")
     private String unit;
-    private long timestamp;
+
+    @SerializedName("Dimensions")
+    private List<Dimension> dimensions;
+
+    @SerializedName("Timestamp")
+    private String timestamp;
+
+    @SerializedName("HostName")
+    private String hostName;
 
     /**
      * Constructs a {@code Metric} instance with the specified {@code metricName} and <code>
@@ -32,7 +59,7 @@ public class Metric {
      * @param value the metric value
      */
     public Metric(String metricName, Number value) {
-        this(metricName, value, "count");
+        this(metricName, value, Unit.COUNT);
     }
 
     /**
@@ -43,11 +70,29 @@ public class Metric {
      * @param value the metric value
      * @param unit the metric unit
      */
-    public Metric(String metricName, Number value, String unit) {
+    public Metric(String metricName, Number value, Unit unit) {
+        this.metricName = metricName;
+        this.value = value.toString();
+        this.unit = unit.getValue();
+        this.hostName = LOCALHOST;
+    }
+
+    /**
+     * Constructs a new {@code Metric} instance.
+     *
+     * @param metricName the metric name
+     * @param unit the metric unit
+     * @param value the metric value
+     * @param hostName the host name
+     * @param timestamp the metric timestamp
+     */
+    private Metric(
+            String metricName, String unit, String value, String hostName, String timestamp) {
         this.metricName = metricName;
-        this.value = value;
         this.unit = unit;
-        timestamp = System.currentTimeMillis();
+        this.value = value;
+        this.hostName = hostName;
+        this.timestamp = timestamp;
     }
 
     /**
@@ -60,12 +105,12 @@ public class Metric {
     }
 
     /**
-     * Returns the value of the {@code Metric}.
+     * Returns the int value of the {@code Metric}.
      *
-     * @return the metric value
+     * @return the metric value in int
      */
-    public Number getValue() {
-        return value;
+    public Double getValue() {
+        return Double.valueOf(value);
     }
 
     /**
@@ -73,8 +118,8 @@ public class Metric {
      *
      * @return the metric unit
      */
-    public String getUnit() {
-        return unit;
+    public Unit getUnit() {
+        return Unit.fromValue(unit);
     }
 
     /**
@@ -82,13 +127,82 @@ public class Metric {
      *
      * @return the metric timestamp
      */
-    public long getTimestamp() {
+    public String getTimestamp() {
         return timestamp;
     }
 
+    /**
+     * Returns the timestamp of the {@code Metric}.
+     *
+     * @return the metric timestamp
+     */
+    public String getHostName() {
+        return hostName;
+    }
+
+    /**
+     * Returns a {@code Metric} instance parsed from the log string.
+     *
+     * @param line the input string
+     * @return a {@code Metric} object
+     */
+    public static Metric parse(String line) {
+        // DiskAvailable.Gigabytes:311|#Level:Host|#hostname:localhost,1650953744320,request_id
+        Matcher matcher = PATTERN.matcher(line);
+        if (!matcher.matches()) {
+            return null;
+        }
+
+        Metric metric =
+                new Metric(
+                        matcher.group(1),
+                        matcher.group(2),
+                        matcher.group(3),
+                        matcher.group(5),
+                        matcher.group(6));
+
+        String dimensions = matcher.group(4);
+        if (dimensions != null) {
+            String[] dimension = dimensions.split(",");
+            List<Dimension> list = new ArrayList<>(dimension.length);
+            for (String dime : dimension) {
+                String[] pair = dime.split(":");
+                if (pair.length == 2) {
+                    list.add(new Dimension(pair[0], pair[1]));
+                }
+            }
+            metric.dimensions = list;
+        }
+
+        return metric;
+    }
+
     /** {@inheritDoc} */
     @Override
     public String toString() {
-        return metricName + '.' + unit + ':' + value + "|#timestamp:" + timestamp;
+        StringBuilder sb = new StringBuilder(128);
+        sb.append(metricName).append('.').append(unit).append(':').append(value).append("|#");
+        if (dimensions != null) {
+            boolean first = true;
+            for (Dimension dimension : dimensions) {
+                if (first) {
+                    first = false;
+                } else {
+                    sb.append(',');
+                }
+                sb.append(dimension.getName()).append(':').append(dimension.getValue());
+            }
+        }
+        sb.append("|#hostname:").append(hostName);
+        sb.append(',').append(timestamp);
+        return sb.toString();
+    }
+
+    private static String getLocalHostName() {
+        try {
+            return InetAddress.getLocalHost().getHostName();
+        } catch (UnknownHostException e) {
+            return "Unknown";
+        }
     }
 }
diff --git a/api/src/main/java/ai/djl/metric/Metrics.java b/api/src/main/java/ai/djl/metric/Metrics.java
index 82ae68400..9412a1632 100644
--- a/api/src/main/java/ai/djl/metric/Metrics.java
+++ b/api/src/main/java/ai/djl/metric/Metrics.java
@@ -12,7 +12,6 @@
  */
 package ai.djl.metric;
 
-import java.io.Serializable;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.Comparator;
@@ -41,8 +40,6 @@ import java.util.stream.Collectors;
  */
 public class Metrics {
 
-    private static final MetricValueComparator VALUE_COMPARATOR = new MetricValueComparator();
-
     private Map<String, List<Metric>> metrics;
 
     /** Constructs an empty {@code Metrics} instance. */
@@ -80,7 +77,7 @@ public class Metrics {
      * @param value the metric value
      * @param unit the metric unit
      */
-    public void addMetric(String name, Number value, String unit) {
+    public void addMetric(String name, Number value, Unit unit) {
         addMetric(new Metric(name, value, unit));
     }
 
@@ -146,7 +143,7 @@ public class Metrics {
         }
 
         List<Metric> list = new ArrayList<>(metric);
-        list.sort(VALUE_COMPARATOR);
+        list.sort(Comparator.comparingDouble(Metric::getValue));
         int index = metric.size() * percentile / 100;
         return list.get(index);
     }
@@ -163,23 +160,6 @@ public class Metrics {
             throw new IllegalArgumentException("Metric name not found: " + metricName);
         }
 
-        return metric.stream().collect(Collectors.averagingDouble(m -> m.getValue().doubleValue()));
-    }
-
-    /** Comparator based on {@code Metric}'s value field. */
-    private static final class MetricValueComparator implements Comparator<Metric>, Serializable {
-
-        private static final long serialVersionUID = 1L;
-
-        /** {@inheritDoc} */
-        @Override
-        public int compare(Metric o1, Metric o2) {
-            Number n1 = o1.getValue();
-            Number n2 = o2.getValue();
-            if (n1 instanceof Double || n1 instanceof Float) {
-                return Double.compare(n1.doubleValue(), n2.doubleValue());
-            }
-            return Long.compare(n1.longValue(), n2.longValue());
-        }
+        return metric.stream().collect(Collectors.averagingDouble(Metric::getValue));
     }
 }
diff --git a/api/src/main/java/ai/djl/metric/Unit.java b/api/src/main/java/ai/djl/metric/Unit.java
new file mode 100644
index 000000000..81d45e631
--- /dev/null
+++ b/api/src/main/java/ai/djl/metric/Unit.java
@@ -0,0 +1,82 @@
+/*
+ * Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.metric;
+
+import java.util.concurrent.ConcurrentHashMap;
+
+/** An interface holds metric unit constants. */
+public enum Unit {
+    MICROSECONDS("Microseconds"),
+    MILLISECONDS("Milliseconds"),
+    BYTES("Bytes"),
+    KILOBYTES("Kilobytes"),
+    MEGABYTES("Megabytes"),
+    GIGABYTES("Gigabytes"),
+    TERABYTES("Terabytes"),
+    BITS("Bits"),
+    KILOBITS("Kilobits"),
+    MEGABITS("Megabits"),
+    GIGABITS("Gigabits"),
+    TERABITS("Terabits"),
+    PERCENT("Percent"),
+    COUNT("Count"),
+    BYTES_PER_SECOND("Bytes/Second"),
+    KILOBYTES_PER_SECOND("Kilobytes/Second"),
+    MEGABYTES_PER_SECOND("Megabytes/Second"),
+    GIGABYTES_PER_SECOND("Gigabytes/Second"),
+    TERABYTES_PER_SECOND("Terabytes/Second"),
+    BITS_PER_SECOND("Bits/Second"),
+    KILOBITS_PER_SECOND("Kilobits/Second"),
+    MEGABITS_PER_SECOND("Megabits/Second"),
+    GIGABITS_PER_SECOND("Gigabits/Second"),
+    TERABITS_PER_SECOND("Terabits/Second"),
+    COUNT_PER_SECOND("Count/Second"),
+    NONE("None");
+
+    private static final ConcurrentHashMap<String, Unit> MAP = new ConcurrentHashMap<>();
+
+    static {
+        for (Unit unit : values()) {
+            MAP.put(unit.value, unit);
+        }
+    }
+
+    private final String value;
+
+    Unit(String value) {
+        this.value = value;
+    }
+
+    /**
+     * Returns the string value of the {@code Unit}.
+     *
+     * @return the string value of the {@code Unit}
+     */
+    public String getValue() {
+        return value;
+    }
+
+    /**
+     * Returns {@code Unit} instance from an string value.
+     *
+     * @param value the String value of Unit
+     * @return the {@code Unit}
+     */
+    public static Unit fromValue(String value) {
+        Unit ret = MAP.get(value);
+        if (ret == null) {
+            throw new IllegalArgumentException("Invalid unit value: " + value);
+        }
+        return ret;
+    }
+}
diff --git a/api/src/main/java/ai/djl/training/listener/MemoryTrainingListener.java b/api/src/main/java/ai/djl/training/listener/MemoryTrainingListener.java
index 0029bb631..fa4b98096 100644
--- a/api/src/main/java/ai/djl/training/listener/MemoryTrainingListener.java
+++ b/api/src/main/java/ai/djl/training/listener/MemoryTrainingListener.java
@@ -15,6 +15,7 @@ package ai.djl.training.listener;
 import ai.djl.Device;
 import ai.djl.metric.Metric;
 import ai.djl.metric.Metrics;
+import ai.djl.metric.Unit;
 import ai.djl.training.Trainer;
 import ai.djl.util.cuda.CudaUtils;
 import java.io.BufferedWriter;
@@ -101,8 +102,8 @@ public class MemoryTrainingListener extends TrainingListenerAdapter {
             long nonHeapUsed = nonHeap.getUsed();
             getProcessInfo(metrics);
 
-            metrics.addMetric("Heap", heapUsed, "bytes");
-            metrics.addMetric("NonHeap", nonHeapUsed, "bytes");
+            metrics.addMetric("Heap", heapUsed, Unit.BYTES);
+            metrics.addMetric("NonHeap", nonHeapUsed, Unit.BYTES);
             int gpuCount = CudaUtils.getGpuCount();
 
             // CudaUtils.getGpuMemory() will allocates memory on GPUs if CUDA runtime is not
@@ -110,7 +111,7 @@ public class MemoryTrainingListener extends TrainingListenerAdapter {
             for (int i = 0; i < gpuCount; ++i) {
                 Device device = Device.gpu(i);
                 MemoryUsage mem = CudaUtils.getGpuMemory(device);
-                metrics.addMetric("GPU-" + i, mem.getCommitted(), "bytes");
+                metrics.addMetric("GPU-" + i, mem.getCommitted(), Unit.BYTES);
             }
         }
     }
@@ -170,8 +171,8 @@ public class MemoryTrainingListener extends TrainingListenerAdapter {
                     }
                     float cpu = Float.parseFloat(tokens[0]);
                     long rss = Long.parseLong(tokens[1]) * 1024;
-                    metrics.addMetric("cpu", cpu, "%");
-                    metrics.addMetric("rss", rss, "bytes");
+                    metrics.addMetric("cpu", cpu, Unit.PERCENT);
+                    metrics.addMetric("rss", rss, Unit.BYTES);
                 }
             } catch (IOException e) {
                 logger.error("Failed execute cmd: " + cmd, e);
diff --git a/engines/onnxruntime/onnxruntime-engine/README.md b/engines/onnxruntime/onnxruntime-engine/README.md
index 87e690fb6..900f987c9 100644
--- a/engines/onnxruntime/onnxruntime-engine/README.md
+++ b/engines/onnxruntime/onnxruntime-engine/README.md
@@ -85,3 +85,11 @@ Gradle:
     }
     implementation "com.microsoft.onnxruntime:onnxruntime_gpu:1.11.0"
 ```
+
+#### Enable TensorRT execution
+
+ONNXRuntime offers TensorRT execution as the backend. In DJL, user can specify the followings in the Criteria to enable:
+
+```
+optOption("ortDevice", "TensorRT")
+```
diff --git a/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/engine/OrtModel.java b/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/engine/OrtModel.java
index dfd9caa1f..7ac6e0ce8 100644
--- a/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/engine/OrtModel.java
+++ b/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/engine/OrtModel.java
@@ -77,10 +77,6 @@ public class OrtModel extends BaseModel {
 
         try {
             SessionOptions ortOptions = getSessionOptions(options);
-            Device device = manager.getDevice();
-            if (device.isGpu()) {
-                ortOptions.addCUDA(manager.getDevice().getDeviceId());
-            }
             OrtSession session = env.createSession(modelFile.toString(), ortOptions);
             block = new OrtSymbolBlock(session, (OrtNDManager) manager);
         } catch (OrtException e) {
@@ -100,10 +96,6 @@ public class OrtModel extends BaseModel {
         try {
             byte[] buf = Utils.toByteArray(is);
             SessionOptions ortOptions = getSessionOptions(options);
-            Device device = manager.getDevice();
-            if (device.isGpu()) {
-                ortOptions.addCUDA(manager.getDevice().getDeviceId());
-            }
             OrtSession session = env.createSession(buf, ortOptions);
             block = new OrtSymbolBlock(session, (OrtNDManager) manager);
         } catch (OrtException e) {
@@ -186,6 +178,28 @@ public class OrtModel extends BaseModel {
             ortSession.setCPUArenaAllocator(true);
         }
 
+        Device device = manager.getDevice();
+        if (options.containsKey("ortDevice")) {
+            String ortDevice = (String) options.get("ortDevice");
+            switch (ortDevice) {
+                case "TensorRT":
+                    if (!device.isGpu()) {
+                        throw new IllegalArgumentException("TensorRT required GPU device.");
+                    }
+                    ortSession.addTensorrt(device.getDeviceId());
+                    break;
+                case "ROCM":
+                    ortSession.addROCM();
+                    break;
+                case "CoreML":
+                    ortSession.addCoreML();
+                    break;
+                default:
+                    throw new IllegalArgumentException("Invalid ortDevice: " + ortDevice);
+            }
+        } else if (device.isGpu()) {
+            ortSession.addCUDA(device.getDeviceId());
+        }
         return ortSession;
     }
 }
diff --git a/extensions/benchmark/src/main/java/ai/djl/benchmark/AbstractBenchmark.java b/extensions/benchmark/src/main/java/ai/djl/benchmark/AbstractBenchmark.java
index 23f7bf5b3..22e413ecf 100644
--- a/extensions/benchmark/src/main/java/ai/djl/benchmark/AbstractBenchmark.java
+++ b/extensions/benchmark/src/main/java/ai/djl/benchmark/AbstractBenchmark.java
@@ -16,6 +16,7 @@ import ai.djl.Device;
 import ai.djl.ModelException;
 import ai.djl.engine.Engine;
 import ai.djl.metric.Metrics;
+import ai.djl.metric.Unit;
 import ai.djl.ndarray.NDList;
 import ai.djl.ndarray.types.DataType;
 import ai.djl.ndarray.types.Shape;
@@ -149,40 +150,28 @@ public abstract class AbstractBenchmark {
                             metrics.getMetric("LoadModel").get(0).getValue().longValue();
                     logger.info(
                             "Model loading time: {} ms.",
-                            String.format("%.3f", loadModelTime / 1_000_000f));
+                            String.format("%.3f", loadModelTime / 1000f));
                 }
 
                 if (metrics.hasMetric("Inference") && iteration > 1) {
-                    float totalP50 =
-                            metrics.percentile("Total", 50).getValue().longValue() / 1_000_000f;
-                    float totalP90 =
-                            metrics.percentile("Total", 90).getValue().longValue() / 1_000_000f;
-                    float totalP99 =
-                            metrics.percentile("Total", 99).getValue().longValue() / 1_000_000f;
-                    float p50 =
-                            metrics.percentile("Inference", 50).getValue().longValue() / 1_000_000f;
-                    float p90 =
-                            metrics.percentile("Inference", 90).getValue().longValue() / 1_000_000f;
-                    float p99 =
-                            metrics.percentile("Inference", 99).getValue().longValue() / 1_000_000f;
+                    float totalP50 = metrics.percentile("Total", 50).getValue().longValue() / 1000f;
+                    float totalP90 = metrics.percentile("Total", 90).getValue().longValue() / 1000f;
+                    float totalP99 = metrics.percentile("Total", 99).getValue().longValue() / 1000f;
+                    float p50 = metrics.percentile("Inference", 50).getValue().longValue() / 1000f;
+                    float p90 = metrics.percentile("Inference", 90).getValue().longValue() / 1000f;
+                    float p99 = metrics.percentile("Inference", 99).getValue().longValue() / 1000f;
                     float preP50 =
-                            metrics.percentile("Preprocess", 50).getValue().longValue()
-                                    / 1_000_000f;
+                            metrics.percentile("Preprocess", 50).getValue().longValue() / 1000f;
                     float preP90 =
-                            metrics.percentile("Preprocess", 90).getValue().longValue()
-                                    / 1_000_000f;
+                            metrics.percentile("Preprocess", 90).getValue().longValue() / 1000f;
                     float preP99 =
-                            metrics.percentile("Preprocess", 99).getValue().longValue()
-                                    / 1_000_000f;
+                            metrics.percentile("Preprocess", 99).getValue().longValue() / 1000f;
                     float postP50 =
-                            metrics.percentile("Postprocess", 50).getValue().longValue()
-                                    / 1_000_000f;
+                            metrics.percentile("Postprocess", 50).getValue().longValue() / 1000f;
                     float postP90 =
-                            metrics.percentile("Postprocess", 90).getValue().longValue()
-                                    / 1_000_000f;
+                            metrics.percentile("Postprocess", 90).getValue().longValue() / 1000f;
                     float postP99 =
-                            metrics.percentile("Postprocess", 99).getValue().longValue()
-                                    / 1_000_000f;
+                            metrics.percentile("Postprocess", 99).getValue().longValue() / 1000f;
                     logger.info(
                             String.format(
                                     "total P50: %.3f ms, P90: %.3f ms, P99: %.3f ms",
@@ -266,12 +255,12 @@ public abstract class AbstractBenchmark {
 
         ZooModel<Void, float[]> model = criteria.loadModel();
         if (device == Device.cpu() || device == Device.gpu()) {
-            long delta = System.nanoTime() - begin;
+            long delta = (System.nanoTime() - begin) / 1000;
             logger.info(
                     "Model {} loaded in: {} ms.",
                     model.getName(),
-                    String.format("%.3f", delta / 1_000_000f));
-            metrics.addMetric("LoadModel", delta);
+                    String.format("%.3f", delta / 1000f));
+            metrics.addMetric("LoadModel", delta, Unit.MICROSECONDS);
         }
         return model;
     }
diff --git a/extensions/benchmark/src/main/java/ai/djl/benchmark/Benchmark.java b/extensions/benchmark/src/main/java/ai/djl/benchmark/Benchmark.java
index a530d2686..25432c75b 100644
--- a/extensions/benchmark/src/main/java/ai/djl/benchmark/Benchmark.java
+++ b/extensions/benchmark/src/main/java/ai/djl/benchmark/Benchmark.java
@@ -18,6 +18,7 @@ import ai.djl.engine.Engine;
 import ai.djl.engine.EngineException;
 import ai.djl.inference.Predictor;
 import ai.djl.metric.Metrics;
+import ai.djl.metric.Unit;
 import ai.djl.repository.zoo.ZooModel;
 import ai.djl.training.listener.MemoryTrainingListener;
 import ai.djl.translate.TranslateException;
@@ -75,14 +76,14 @@ public final class Benchmark extends AbstractBenchmark {
                 predictor.predict(null); // warmup
 
                 predictor.setMetrics(metrics); // Let predictor collect metrics
-                metrics.addMetric("start", System.currentTimeMillis(), "mills");
+                metrics.addMetric("start", System.currentTimeMillis(), Unit.MILLISECONDS);
                 for (int i = 0; i < iteration; ++i) {
                     predictResult = predictor.predict(null);
 
                     progressBar.update(i);
                     MemoryTrainingListener.collectMemoryInfo(metrics);
                 }
-                metrics.addMetric("end", System.currentTimeMillis(), "mills");
+                metrics.addMetric("end", System.currentTimeMillis(), Unit.MILLISECONDS);
             }
             return predictResult;
         }
diff --git a/extensions/benchmark/src/main/java/ai/djl/benchmark/MultithreadedBenchmark.java b/extensions/benchmark/src/main/java/ai/djl/benchmark/MultithreadedBenchmark.java
index 06a64af6f..8c2456152 100644
--- a/extensions/benchmark/src/main/java/ai/djl/benchmark/MultithreadedBenchmark.java
+++ b/extensions/benchmark/src/main/java/ai/djl/benchmark/MultithreadedBenchmark.java
@@ -17,6 +17,7 @@ import ai.djl.ModelException;
 import ai.djl.engine.Engine;
 import ai.djl.inference.Predictor;
 import ai.djl.metric.Metrics;
+import ai.djl.metric.Unit;
 import ai.djl.repository.zoo.ZooModel;
 import ai.djl.training.listener.MemoryTrainingListener;
 import ai.djl.translate.TranslateException;
@@ -83,7 +84,7 @@ public class MultithreadedBenchmark extends AbstractBenchmark {
                 callable.warmup();
             }
 
-            metrics.addMetric("start", System.currentTimeMillis(), "mills");
+            metrics.addMetric("start", System.currentTimeMillis(), Unit.MILLISECONDS);
             try {
                 List<Future<float[]>> futures;
                 if (delay > 0) {
@@ -105,7 +106,7 @@ public class MultithreadedBenchmark extends AbstractBenchmark {
             } catch (InterruptedException | ExecutionException e) {
                 logger.error("", e);
             }
-            metrics.addMetric("end", System.currentTimeMillis(), "mills");
+            metrics.addMetric("end", System.currentTimeMillis(), Unit.MILLISECONDS);
             for (PredictorCallable callable : callables) {
                 callable.close();
             }


Test patch:
diff --git a/api/src/test/java/ai/djl/metric/MetricsTest.java b/api/src/test/java/ai/djl/metric/MetricsTest.java
index 092a5f7f7..6fec58d51 100644
--- a/api/src/test/java/ai/djl/metric/MetricsTest.java
+++ b/api/src/test/java/ai/djl/metric/MetricsTest.java
@@ -23,22 +23,25 @@ public class MetricsTest {
     public void testMetrics() {
         Metrics metrics = new Metrics();
         metrics.addMetric(new Metric("m1", 1L));
-        metrics.addMetric("m1", 3L, "count");
+        metrics.addMetric("m1", 3L, Unit.COUNT);
         metrics.addMetric("m1", 2L);
         Metric p50 = metrics.percentile("m1", 50);
         Assert.assertEquals(p50.getValue().longValue(), 2L);
 
         metrics.addMetric(new Metric("m2", 1f));
-        metrics.addMetric("m2", 3f, "count");
+        metrics.addMetric("m2", 3f, Unit.COUNT);
         metrics.addMetric("m2", 2f);
         p50 = metrics.percentile("m2", 50);
         Assert.assertEquals(p50.getValue().floatValue(), 2f);
 
         metrics.addMetric(new Metric("m3", 1d));
-        metrics.addMetric("m3", 3d, "count");
+        metrics.addMetric("m3", 3d, Unit.COUNT);
         metrics.addMetric("m3", 2d);
         p50 = metrics.percentile("m3", 50);
         Assert.assertEquals(p50.getValue().doubleValue(), 2d);
+        Assert.assertTrue(metrics.hasMetric("m1"));
+        Assert.assertEquals(2L, metrics.latestMetric("m1").getValue().longValue());
+        Assert.assertThrows(() -> metrics.latestMetric("none"));
 
         List<Metric> list = metrics.getMetric("m1");
         Assert.assertEquals(list.size(), 3);
@@ -54,6 +57,23 @@ public class MetricsTest {
         Assert.assertFalse(metricNames.contains("m4"));
     }
 
+    @Test
+    public void testParseMetrics() {
+        String line = "Disk.Gigabytes:311|#Level:Host|#hostname:localhost,1650953744320";
+        Metric metric = Metric.parse(line);
+        Assert.assertNotNull(metric);
+        Assert.assertEquals(metric.getValue().intValue(), 311);
+        Assert.assertEquals(metric.getTimestamp(), "1650953744320");
+        Assert.assertEquals(metric.getUnit(), Unit.GIGABYTES);
+        Assert.assertEquals(metric.getHostName(), "localhost");
+        Assert.assertEquals(line, metric.toString());
+
+        Assert.assertNull(Metric.parse("DiskAvailable.Gigabytes:311"));
+        Metric invalid = Metric.parse("Disk.InvalidUnits:311|#Level:Host|#hostname:localhost,1,X");
+        Assert.assertNotNull(invalid);
+        Assert.expectThrows(IllegalArgumentException.class, invalid::getUnit);
+    }
+
     @Test
     public void testMetricsMean() {
         Metrics metrics = new Metrics();