Removing softmax as terminal activation within remaining image classi…

…fication models. (#491)
tensorflow · May 6, 2020 · b0cc1c4 · b0cc1c4
1 parent 08692fb
commit b0cc1c4
Show file tree

Hide file tree

Showing 5 changed files with 7 additions and 8 deletions.
diff --git a/Models/ImageClassification/EfficientNet.swift b/Models/ImageClassification/EfficientNet.swift
@@ -271,7 +271,7 @@ public struct EfficientNet: Layer {
         dropoutProb = Dropout<Float>(probability: dropout)
         outputClassifier = Dense(
             inputSize: makeDivisible(filter: 1280, width: width),
-            outputSize: classCount, activation: softmax)
+            outputSize: classCount)
     }
 
     @differentiable

diff --git a/Models/ImageClassification/MobileNetV2.swift b/Models/ImageClassification/MobileNetV2.swift
@@ -219,8 +219,7 @@ public struct MobileNetV2: Layer {
         outputConvBatchNorm = BatchNorm(featureCount: lastBlockFilterCount)
 
         outputClassifier = Dense(
-            inputSize: lastBlockFilterCount, outputSize: classCount,
-            activation: softmax)
+            inputSize: lastBlockFilterCount, outputSize: classCount)
     }
 
     @differentiable

diff --git a/Models/ImageClassification/MobileNetV3.swift b/Models/ImageClassification/MobileNetV3.swift
@@ -357,7 +357,7 @@ public struct MobileNetV3Large: Layer {
             input.shape[0], 1, 1, self.lastConvChannel,
         ])
         let finalConvResult = dropoutLayer(hardSwish(finalConv(averagePool)))
-        return softmax(flatten(classiferConv(finalConvResult)))
+        return flatten(classiferConv(finalConvResult))
     }
 }
 
@@ -473,6 +473,6 @@ public struct MobileNetV3Small: Layer {
             input.shape[0], 1, 1, lastConvChannel,
         ])
         let finalConvResult = dropoutLayer(hardSwish(finalConv(averagePool)))
-        return softmax(flatten(classiferConv(finalConvResult)))
+        return flatten(classiferConv(finalConvResult))
     }
 }
diff --git a/Models/ImageClassification/VGG.swift b/Models/ImageClassification/VGG.swift
@@ -58,7 +58,7 @@ public struct VGG16: Layer {
         layer3 = VGGBlock(featureCounts: (128, 256, 256, 256), blockCount: 3)
         layer4 = VGGBlock(featureCounts: (256, 512, 512, 512), blockCount: 3)
         layer5 = VGGBlock(featureCounts: (512, 512, 512, 512), blockCount: 3)
-        output = Dense(inputSize: 4096, outputSize: classCount, activation: softmax)
+        output = Dense(inputSize: 4096, outputSize: classCount)
     }
 
     @differentiable
@@ -86,7 +86,7 @@ public struct VGG19: Layer {
         layer3 = VGGBlock(featureCounts: (128, 256, 256, 256), blockCount: 4)
         layer4 = VGGBlock(featureCounts: (256, 512, 512, 512), blockCount: 4)
         layer5 = VGGBlock(featureCounts: (512, 512, 512, 512), blockCount: 4)
-        output = Dense(inputSize: 4096, outputSize: classCount, activation: softmax)
+        output = Dense(inputSize: 4096, outputSize: classCount)
     }
 
     @differentiable

diff --git a/Models/Spatiotemporal/C3D.swift b/Models/Spatiotemporal/C3D.swift
@@ -39,7 +39,7 @@ public struct C3D: Layer {
     var output: Dense<Float>
 
     public init(classCount: Int) {
-        self.output = Dense<Float>(inputSize: 1024, outputSize: classCount, activation: softmax)
+        self.output = Dense<Float>(inputSize: 1024, outputSize: classCount)
     }
 
     @differentiable