Repository: deeplearning4j/deeplearning4j
Issue: Unet: Error importing model for training with Conv2D as last layer (0.9.2-SNAPSHOT)

#### Issue Description

Hi so I have a model in keras that looks like this. 
 
``` python
def get_unet():
    conv_params = dict(activation='relu', border_mode='same')
    merge_params = dict(mode='concat', concat_axis=1)
    inputs = Input((8, 256, 256))
    conv1 = Convolution2D(32, 3, 3, **conv_params)(inputs)
    conv1 = Convolution2D(32, 3, 3, **conv_params)(conv1)
    pool1 = MaxPooling2D(pool_size=(2, 2))(conv1)

    conv2 = Convolution2D(64, 3, 3, **conv_params)(pool1)
    conv2 = Convolution2D(64, 3, 3, **conv_params)(conv2)
    pool2 = MaxPooling2D(pool_size=(2, 2))(conv2)

    conv3 = Convolution2D(128, 3, 3, **conv_params)(pool2)
    conv3 = Convolution2D(128, 3, 3, **conv_params)(conv3)
    pool3 = MaxPooling2D(pool_size=(2, 2))(conv3)

    conv4 = Convolution2D(256, 3, 3, **conv_params)(pool3)
    conv4 = Convolution2D(256, 3, 3, **conv_params)(conv4)
    pool4 = MaxPooling2D(pool_size=(2, 2))(conv4)

    conv5 = Convolution2D(512, 3, 3, **conv_params)(pool4)
    conv5 = Convolution2D(512, 3, 3, **conv_params)(conv5)

    up6 = merge_l([UpSampling2D(size=(2, 2))(conv5), conv4], **merge_params)
    conv6 = Convolution2D(256, 3, 3, **conv_params)(up6)
    conv6 = Convolution2D(256, 3, 3, **conv_params)(conv6)

    up7 = merge_l([UpSampling2D(size=(2, 2))(conv6), conv3], **merge_params)
    conv7 = Convolution2D(128, 3, 3, **conv_params)(up7)
    conv7 = Convolution2D(128, 3, 3, **conv_params)(conv7)

    up8 = merge_l([UpSampling2D(size=(2, 2))(conv7), conv2], **merge_params)
    conv8 = Convolution2D(64, 3, 3, **conv_params)(up8)
    conv8 = Convolution2D(64, 3, 3, **conv_params)(conv8)

    up9 = merge_l([UpSampling2D(size=(2, 2))(conv8), conv1], **merge_params)
    conv9 = Convolution2D(32, 3, 3, **conv_params)(up9)
    conv9 = Convolution2D(32, 3, 3, **conv_params)(conv9)

    conv10 = Convolution2D(1, 1, 1, activation='sigmoid')(conv9)

    optimizer = SGD(lr=0.01, momentum=0.9, nesterov=True)
    model = Model(input=inputs, output=conv10)
    model.compile(optimizer=optimizer,
                  loss='binary_crossentropy',
                  metrics=['accuracy', jaccard_coef, jaccard_coef_int])
    return model
```
And saved like this 
``` python
model = get_unet()

#train model on lots of images....

model.save('unet.h5')
```
I would like to try loading this trained model into dl4j. I tried to do this with 
``` scala
  val unet = KerasModelImport.importKerasModelAndWeights("src/main/resources/unet.h5")
```  
However I got the following error:
```
Exception in thread "main" java.lang.IllegalStateException: Invalid input type (layer index = -1, layer name="convolution2d_19_loss"): expected FeedForward input type. Got: InputTypeConvolutional(h=256,w=256,d=1)

```
Any ideas? This model works as expected in keras. I had to use the snapshot build so that the merge layers will work in dl4j. Not sure why it doesn't like that the first layer is convolutional. 

keras==1.0.0
theano==0.9.0
inside keras.json 
```
{
    "image_dim_ordering": "th",
    "epsilon": 1e-07,
    "floatx": "float32",
    "backend": "theano"
}
```

Changed files:
deeplearning4j-modelimport/src/main/java/org/deeplearning4j/nn/modelimport/keras/KerasModel.java
deeplearning4j-modelimport/src/test/java/org/deeplearning4j/nn/modelimport/keras/e2e/KerasModelEndToEndTest.java
deeplearning4j-modelimport/src/main/java/org/deeplearning4j/nn/modelimport/keras/layers/KerasLoss.java

diff --git a/deeplearning4j-modelimport/src/main/java/org/deeplearning4j/nn/modelimport/keras/KerasModel.java b/deeplearning4j-modelimport/src/main/java/org/deeplearning4j/nn/modelimport/keras/KerasModel.java
index c11dd47493..1b23683f2f 100644
--- a/deeplearning4j-modelimport/src/main/java/org/deeplearning4j/nn/modelimport/keras/KerasModel.java
+++ b/deeplearning4j-modelimport/src/main/java/org/deeplearning4j/nn/modelimport/keras/KerasModel.java
@@ -322,17 +322,10 @@ public class KerasModel {
                 if (preprocessor != null)
                     preprocessors.put(layer.getLayerName(), preprocessor);
                 graphBuilder.addLayer(layer.getLayerName(), layer.getLayer(), inboundLayerNamesArray);
-                if (this.outputLayerNames.contains(layer.getLayerName()) && !(layer.getLayer() instanceof IOutputLayer))
-                    log.warn("Model cannot be trained: output layer " + layer.getLayerName()
-                            + " is not an IOutputLayer (no loss function specified)");
             } else if (layer.isVertex()) { // Ignore "preprocessor" layers for now
                 if (preprocessor != null)
                     preprocessors.put(layer.getLayerName(), preprocessor);
                 graphBuilder.addVertex(layer.getLayerName(), layer.getVertex(), inboundLayerNamesArray);
-                if (this.outputLayerNames.contains(layer.getLayerName())
-                        && !(layer.getVertex() instanceof IOutputLayer))
-                    log.warn("Model cannot be trained: output vertex " + layer.getLayerName()
-                            + " is not an IOutputLayer (no loss function specified)");
             } else if (layer.isInputPreProcessor()) {
                 if (preprocessor == null)
                     throw new UnsupportedKerasConfigurationException("Layer " + layer.getLayerName()
@@ -340,10 +333,6 @@ public class KerasModel {
                 graphBuilder.addVertex(layer.getLayerName(), new PreprocessorVertex(preprocessor),
                         inboundLayerNamesArray);
             }
-
-            if (this.outputLayerNames.contains(layer.getLayerName()))
-                log.warn("Model cannot be trained: output " + layer.getLayerName()
-                        + " is not an IOutputLayer (no loss function specified)");
         }
         graphBuilder.setInputPreProcessors(preprocessors);
 
diff --git a/deeplearning4j-modelimport/src/main/java/org/deeplearning4j/nn/modelimport/keras/layers/KerasLoss.java b/deeplearning4j-modelimport/src/main/java/org/deeplearning4j/nn/modelimport/keras/layers/KerasLoss.java
index 33884e2121..4e264de468 100644
--- a/deeplearning4j-modelimport/src/main/java/org/deeplearning4j/nn/modelimport/keras/layers/KerasLoss.java
+++ b/deeplearning4j-modelimport/src/main/java/org/deeplearning4j/nn/modelimport/keras/layers/KerasLoss.java
@@ -4,7 +4,10 @@ import lombok.Data;
 import lombok.EqualsAndHashCode;
 import lombok.extern.slf4j.Slf4j;
 import org.deeplearning4j.nn.conf.inputs.InputType;
+import org.deeplearning4j.nn.conf.layers.CnnLossLayer;
+import org.deeplearning4j.nn.conf.layers.FeedForwardLayer;
 import org.deeplearning4j.nn.conf.layers.LossLayer;
+import org.deeplearning4j.nn.conf.layers.RnnLossLayer;
 import org.deeplearning4j.nn.modelimport.keras.KerasLayer;
 import org.deeplearning4j.nn.modelimport.keras.exceptions.InvalidKerasConfigurationException;
 import org.deeplearning4j.nn.modelimport.keras.exceptions.UnsupportedKerasConfigurationException;
@@ -25,6 +28,8 @@ import static org.deeplearning4j.nn.modelimport.keras.utils.KerasLossUtils.mapLo
 public class KerasLoss extends KerasLayer {
 
     private final String KERAS_CLASS_NAME_LOSS = "Loss";
+    private LossFunctions.LossFunction loss;
+
 
     /**
      * Constructor from layer name and input shape.
@@ -58,7 +63,6 @@ public class KerasLoss extends KerasLayer {
         this.dimOrder = DimOrder.NONE;
         this.inboundLayerNames = new ArrayList<>();
         this.inboundLayerNames.add(inboundLayerName);
-        LossFunctions.LossFunction loss;
         try {
             loss = mapLossFunction(kerasLoss, conf);
         } catch (UnsupportedKerasConfigurationException e) {
@@ -67,7 +71,6 @@ public class KerasLoss extends KerasLayer {
             log.warn("Unsupported Keras loss function. Replacing with MSE.");
             loss = LossFunctions.LossFunction.SQUARED_LOSS;
         }
-        this.layer = new LossLayer.Builder(loss).name(layerName).build();
     }
 
     /**
@@ -75,8 +78,20 @@ public class KerasLoss extends KerasLayer {
      *
      * @return LossLayer
      */
-    public LossLayer getLossLayer() {
-        return (LossLayer) this.layer;
+    public FeedForwardLayer getLossLayer(InputType type) throws UnsupportedKerasConfigurationException {
+        if (type instanceof InputType.InputTypeFeedForward) {
+            this.layer = new LossLayer.Builder(loss).name(this.layerName).build();
+        }
+        else if (type instanceof  InputType.InputTypeRecurrent) {
+            this.layer = new RnnLossLayer.Builder(loss).name(this.layerName).build();
+        }
+        else if (type instanceof InputType.InputTypeConvolutional) {
+            this.layer = new CnnLossLayer.Builder(loss).name(this.layerName).build();
+        } else {
+            throw new UnsupportedKerasConfigurationException("Unsupported output layer type"
+                    + "got : " + type.toString());
+        }
+        return (FeedForwardLayer) this.layer;
     }
 
     /**
@@ -87,10 +102,11 @@ public class KerasLoss extends KerasLayer {
      * @throws InvalidKerasConfigurationException Invalid Keras config
      */
     @Override
-    public InputType getOutputType(InputType... inputType) throws InvalidKerasConfigurationException {
+    public InputType getOutputType(InputType... inputType) throws InvalidKerasConfigurationException,
+    UnsupportedKerasConfigurationException {
         if (inputType.length > 1)
             throw new InvalidKerasConfigurationException(
                     "Keras Loss layer accepts only one input (received " + inputType.length + ")");
-        return this.getLossLayer().getOutputType(-1, inputType[0]);
+        return this.getLossLayer(inputType[0]).getOutputType(-1, inputType[0]);
     }
 }
diff --git a/deeplearning4j-modelimport/src/test/java/org/deeplearning4j/nn/modelimport/keras/e2e/KerasModelEndToEndTest.java b/deeplearning4j-modelimport/src/test/java/org/deeplearning4j/nn/modelimport/keras/e2e/KerasModelEndToEndTest.java
index 68c8f89e49..8d0cf1d141 100644
--- a/deeplearning4j-modelimport/src/test/java/org/deeplearning4j/nn/modelimport/keras/e2e/KerasModelEndToEndTest.java
+++ b/deeplearning4j-modelimport/src/test/java/org/deeplearning4j/nn/modelimport/keras/e2e/KerasModelEndToEndTest.java
@@ -287,6 +287,15 @@ public class KerasModelEndToEndTest {
         importFunctionalModelH5Test("modelimport/keras/examples/reshape_to_rnn/reshape_model.h5");
     }
 
+    /**
+     * U-Net
+     */
+    @Test
+    public void importUnetTfKeras2() throws Exception {
+        importFunctionalModelH5Test(
+                "modelimport/keras/examples/unet/unet_keras_2_tf.h5", null, true);
+    }
+
     /**
      * ResNet50
      */
@@ -329,7 +338,8 @@ public class KerasModelEndToEndTest {
     @Ignore
     // Takes unreasonably long, but works
     public void importInception() throws Exception {
-        ComputationGraph graph = importFunctionalModelH5Test("modelimport/keras/examples/inception/inception_v3_complete.h5");
+        ComputationGraph graph = importFunctionalModelH5Test(
+                "modelimport/keras/examples/inception/inception_v3_complete.h5");
         INDArray input = Nd4j.ones(10, 3, 299, 299);
         graph.output(input);
         System.out.println(graph.summary());
@@ -343,23 +353,22 @@ public class KerasModelEndToEndTest {
     public void importXception() throws Exception {
         int[] inputShape = new int[]{299, 299, 3};
         ComputationGraph graph = importFunctionalModelH5Test(
-                "modelimport/keras/examples/xception/xception_tf_keras_2.h5", inputShape);
+                "modelimport/keras/examples/xception/xception_tf_keras_2.h5", inputShape, false);
     }
 
 
-
     private ComputationGraph importFunctionalModelH5Test(String modelPath) throws Exception {
-        return importFunctionalModelH5Test(modelPath, null);
+        return importFunctionalModelH5Test(modelPath, null, false);
     }
 
-    private ComputationGraph importFunctionalModelH5Test(String modelPath, int[] inputShape) throws Exception {
+    private ComputationGraph importFunctionalModelH5Test(String modelPath, int[] inputShape, boolean train) throws Exception {
         ClassPathResource modelResource =
                 new ClassPathResource(modelPath,
                         KerasModelEndToEndTest.class.getClassLoader());
         File modelFile = File.createTempFile(TEMP_MODEL_FILENAME, H5_EXTENSION);
         Files.copy(modelResource.getInputStream(), modelFile.toPath(), StandardCopyOption.REPLACE_EXISTING);
         KerasModelBuilder builder = new KerasModel().modelBuilder().modelHdf5Filename(modelFile.getAbsolutePath())
-                .enforceTrainingConfig(false);
+                .enforceTrainingConfig(train);
         if (inputShape != null) {
             builder.inputShape(inputShape);
         }