title: Translating Pytorch program into Keras: different results
body: <p>I have translated a pytorch program into keras.</p>
  
  <p>A working Pytorch program:</p>
  
  <pre><code>import numpy as np
  import cv2
  import torch
  import torch.nn as nn
  from skimage import segmentation
  
  np.random.seed(1)
  torch.manual_seed(1)
  
  
  fi = "in.jpg"
  
  
  class MyNet(nn.Module):
      def __init__(self, n_inChannel, n_outChannel):
          super(MyNet, self).__init__()
  
          self.seq = nn.Sequential(
              nn.Conv2d(n_inChannel, n_outChannel, kernel_size=3, stride=1, padding=1),
              nn.ReLU(inplace=True),
              nn.BatchNorm2d(n_outChannel),
  
              nn.Conv2d(n_outChannel, n_outChannel, kernel_size=3, stride=1, padding=1),
              nn.ReLU(inplace=True),
              nn.BatchNorm2d(n_outChannel),
  
              nn.Conv2d(n_outChannel, n_outChannel, kernel_size=1, stride=1, padding=0),
              nn.BatchNorm2d(n_outChannel)
          )
  
      def forward(self, x):
          return self.seq(x)
  
  
  im = cv2.imread(fi)
  data = torch.from_numpy(np.array([im.transpose((2, 0, 1)).astype('float32')/255.]))
  data = data.cuda()
  
  labels = segmentation.slic(im, compactness=100, n_segments=10000)
  labels = labels.flatten()
  u_labels = np.unique(labels)
  label_indexes = np.array([np.where(labels == u_label)[0] for u_label in u_labels])
  
  n_inChannel = 3
  n_outChannel = 100
  
  model = MyNet(n_inChannel, n_outChannel)
  model.cuda()
  model.train()
  loss_fn = torch.nn.CrossEntropyLoss()
  optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
  
  label_colours = np.random.randint(255,size=(100,3))
  
  for batch_idx in range(100):
      optimizer.zero_grad()
      output = model( data )[ 0 ]
      output = output.permute( 1, 2, 0 ).view(-1, n_outChannel)
      ignore, target = torch.max( output, 1 )
  
      im_target = target.data.cpu().numpy()
      nLabels = len(np.unique(im_target))
  
      im_target_rgb = np.array([label_colours[ c % 100 ] for c in im_target]) # correct position of "im_target"
      im_target_rgb = im_target_rgb.reshape( im.shape ).astype( np.uint8 )
  
      for inds in label_indexes:
          u_labels_, hist = np.unique(im_target[inds], return_counts=True) 
          im_target[inds] = u_labels_[np.argmax(hist, 0)]
  
      target = torch.from_numpy(im_target)
      target = target.cuda()    
  
      loss = loss_fn(output, target)
  
      loss.backward()
      optimizer.step()
  
      print (batch_idx, '/', 100, ':', nLabels, loss.item())
  
      if nLabels &lt;= 3:
          break         
  fo = "out.jpg"
  cv2.imwrite(fo, im_target_rgb)
  </code></pre>
  
  <p>(source: <a href="https://github.com/kanezaki/pytorch-unsupervised-segmentation/blob/master/demo.py" rel="nofollow noreferrer">https://github.com/kanezaki/pytorch-unsupervised-segmentation/blob/master/demo.py</a>)</p>
  
  <p>My translation into Keras:</p>
  
  <pre><code>import cv2
  import numpy as np
  from skimage import segmentation
  from keras.layers import Conv2D, BatchNormalization, Input, Reshape
  from keras.models import Model
  import keras.backend as k
  from keras.optimizers import SGD, Adam
  from skimage.util import img_as_float
  from skimage import io
  from keras.models import Sequential
  
  np.random.seed(0)
  
  fi = "in.jpg"
  im = cv2.imread(fi).astype(float)/255.
  
  labels = segmentation.slic(im, compactness=100, n_segments=10000)
  labels = labels.flatten()
  print (labels.shape)
  u_labels = np.unique(labels)
  
  label_indexes = [np.where(labels == u_label)[0] for u_label in np.unique(labels)]
  
  n_channels = 100
  
  model = Sequential()
  
  model.add (  Conv2D(n_channels, kernel_size=3, activation='relu', input_shape=im.shape, padding='same'))    
  model.add( BatchNormalization())
  
  model.add( Conv2D(n_channels, kernel_size=3, activation='relu', padding='same'))    
  model.add( BatchNormalization())
  
  model.add( Conv2D(n_channels, kernel_size=1, padding='same'))
  model.add( BatchNormalization())
  
  model.add( Reshape((im.shape[0] * im.shape[1], n_channels)))
  
  img = np.expand_dims(im,0)
  print (img.shape)
  output = model.predict(img)
  
  print (output.shape)
  
  im_target = np.argmax(output[0], 1)
  print (im_target.shape)
  
  for inds in label_indexes:
      u_labels_, hist = np.unique(im_target[inds], return_counts=True)
      im_target[inds] = u_labels_[np.argmax(hist, 0)]
  
  def custom_loss(loss_target, loss_output):
    return k.categorical_crossentropy(target=k.stack(loss_target), output=k.stack(loss_output), from_logits=True)
  
  
  model.compile(optimizer=SGD(lr=0.1, momentum=0.9), loss=custom_loss)
  
  model.fit(img, output, epochs=100, batch_size=1, verbose=1)
  
  pred_result = model.predict(x=[img])[0]
  print (pred_result.shape)
  
  target = np.argmax(pred_result, 1)
  print (target.shape)
  
  nLabels = len(np.unique(target))
  
  label_colours = np.random.randint(255, size=(100, 3))
  
  im_target_rgb = np.array([label_colours[c % 100] for c in im_target])
  
  im_target_rgb = im_target_rgb.reshape(im.shape).astype(np.uint8)
  
  cv2.imwrite("out.jpg", im_target_rgb)
  </code></pre>
  
  <p><em>However, Keras output is really different than of pytorch</em></p>
  
  <p>Input image:</p>
  
  <p><a href="https://i.stack.imgur.com/GHxoA.jpg" rel="nofollow noreferrer"><img src="https://i.stack.imgur.com/GHxoA.jpg" alt="input image"></a></p>
  
  <p>Pytorch result:</p>
  
  <p><a href="https://i.stack.imgur.com/PTRQY.jpg" rel="nofollow noreferrer"><img src="https://i.stack.imgur.com/PTRQY.jpg" alt="pytorch result"></a></p>
  
  <p>Keras result:</p>
  
  <p><a href="https://i.stack.imgur.com/iJLfU.jpg" rel="nofollow noreferrer"><img src="https://i.stack.imgur.com/iJLfU.jpg" alt="keras result"></a></p>
  
  <p>Could someone help me for this translation?  </p>
  
  <p><strong>Edit 1:</strong></p>
  
  <pre><code>I corrected two errors as advised by @sebrockm
  1. removed `relu` from last conv layer
  2. added `from_logits = True` in the loss function
  </code></pre>
  
  <p>Also, changed the no. of conv layers from 4 to 3 to match with the original code. </p>
  
  <pre><code>However, output image did not improve than before and the `loss` are resulted in negative:
  
  Epoch 99/100
  1/1 [==============================] - 0s 92ms/step - loss: -22.8380
  
  Epoch 100/100    
  1/1 [==============================] - 0s 99ms/step - loss: -23.039
  </code></pre>
  
  <p><strong>I think that the Keras code lacks connection between model and output. However, could not figure out to make this connection.</strong></p>