Label -> StringLabel, description of ImageEncoder

Signed-off-by: Tomasz Kornuta <[email protected]>
NVIDIA · Jun 3, 2020 · eee4c1e · eee4c1e
1 parent 6613d40
commit eee4c1e
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 9 deletions.
diff --git a/nemo/collections/cv/modules/data_layers/cifar100_datalayer.py b/nemo/collections/cv/modules/data_layers/cifar100_datalayer.py
@@ -44,7 +44,7 @@
 from torchvision.transforms import Compose, Resize, ToTensor
 
 from nemo.backends.pytorch.nm import DataLayerNM
-from nemo.core.neural_types import AxisKind, AxisType, ClassificationTarget, ImageValue, Index, Label, NeuralType
+from nemo.core.neural_types import AxisKind, AxisType, ClassificationTarget, ImageValue, Index, NeuralType, StringLabel
 from nemo.utils.decorators import add_port_docs
 
 __all__ = ['CIFAR100DataLayer']
@@ -172,9 +172,9 @@ def output_ports(self):
                 elements_type=ImageValue(),  # uint8, <0-255>
             ),
             "coarse_targets": NeuralType(tuple('B'), elements_type=ClassificationTarget()),
-            "coarse_labels": NeuralType(tuple('B'), elements_type=Label()),  # Labels is string!
+            "coarse_labels": NeuralType(tuple('B'), elements_type=StringLabel()),  # Labels is string!
             "fine_targets": NeuralType(tuple('B'), elements_type=ClassificationTarget()),
-            "fine_labels": NeuralType(tuple('B'), elements_type=Label()),  # Labels is string!
+            "fine_labels": NeuralType(tuple('B'), elements_type=StringLabel()),  # Labels is string!
         }
 
     def __len__(self):

diff --git a/nemo/collections/cv/modules/data_layers/mnist_datalayer.py b/nemo/collections/cv/modules/data_layers/mnist_datalayer.py
@@ -28,9 +28,9 @@
     AxisType,
     ClassificationTarget,
     Index,
-    Label,
     NeuralType,
     NormalizedImageValue,
+    StringLabel,
 )
 from nemo.utils.decorators import add_port_docs
 
@@ -110,7 +110,7 @@ def output_ports(self):
                 elements_type=NormalizedImageValue(),  # float, <0-1>
             ),
             "targets": NeuralType(tuple('B'), elements_type=ClassificationTarget()),  # Target are ints!
-            "labels": NeuralType(tuple('B'), elements_type=Label()),  # Labels is string!
+            "labels": NeuralType(tuple('B'), elements_type=StringLabel()),  # Labels is string!
         }
 
     def __len__(self):

diff --git a/nemo/collections/cv/modules/trainables/image_encoder.py b/nemo/collections/cv/modules/trainables/image_encoder.py
@@ -54,7 +54,9 @@
 
 class ImageEncoder(TrainableNM):
     """
-    Class
+    Neural Module implementing a general-usage image encoderds.
+    It encapsulates several models from TorchVision (VGG16, ResNet152 and DensNet121, naming a few).
+    Offers two operation modes and can return: image embeddings vs feature maps.
     """
 
     def __init__(

diff --git a/nemo/core/neural_types/elements.py b/nemo/core/neural_types/elements.py
@@ -41,7 +41,7 @@
     'Index',
     'ImageValue',
     'NormalizedImageValue',
-    'Label',
+    'StringLabel',
 ]
 
 import abc
@@ -219,9 +219,9 @@ class ClassificationTarget(Target):
     """
 
 
-class Label(ElementType):
+class StringLabel(ElementType):
     """
-        Type representing an element being a target value.
+        Type representing an label being a string with class name (e.g. the "hamster" class in CIFAR100).
     """