diff --git a/.gitignore b/.gitignore
index a014da0..025cf91 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,4 +1,6 @@
 *.pyc
 data/
 __pycache__/
+checkpoints/
 *.pth
+
diff --git a/data_vis.py b/data_vis.py
index 714acc5..365e4d1 100644
--- a/data_vis.py
+++ b/data_vis.py
@@ -1,5 +1,6 @@
 import matplotlib.pyplot as plt
 
+
 def plot_img_mask(img, mask):
     fig = plt.figure()
 
diff --git a/load.py b/load.py
index a981f03..c4418c1 100644
--- a/load.py
+++ b/load.py
@@ -1,47 +1,42 @@
+
+#
+# load.py : utils on generators / lists of ids to transform from strings to
+#           cropped images and masks
+
 import os
-import random
 import numpy as np
+
 from PIL import Image
 from functools import partial
-from utils import resize_and_crop, get_square
+from utils import resize_and_crop, get_square, normalize
 
 
 def get_ids(dir):
     """Returns a list of the ids in the directory"""
     return (f[:-4] for f in os.listdir(dir))
 
+
 def split_ids(ids, n=2):
     """Split each id in n, creating n tuples (id, k) for each id"""
     return ((id, i) for i in range(n) for id in ids)
 
-def shuffle_ids(ids):
-    """Returns a shuffle list od the ids"""
-    lst = list(ids)
-    random.shuffle(lst)
-    return lst
 
 def to_cropped_imgs(ids, dir, suffix):
-    """From a list of tuples, returns the correct cropped img (left or right)"""
+    """From a list of tuples, returns the correct cropped img"""
     for id, pos in ids:
         im = resize_and_crop(Image.open(dir + id + suffix))
         yield get_square(im, pos)
 
 
-
-def get_imgs_and_masks():
-    """From the list of ids, return the couples (img, mask)"""
-    dir_img = 'data/train/'
-    dir_mask = 'data/train_masks/'
-
-    ids = get_ids(dir_img)
-    ids = split_ids(ids)
-    ids = shuffle_ids(ids)
+def get_imgs_and_masks(ids, dir_img, dir_mask):
+    """Return all the couples (img, mask)"""
 
     imgs = to_cropped_imgs(ids, dir_img, '.jpg')
 
     # need to transform from HWC to CHW
     imgs_switched = map(partial(np.transpose, axes=[2, 0, 1]), imgs)
+    imgs_normalized = map(normalize, imgs_switched)
 
     masks = to_cropped_imgs(ids, dir_mask, '_mask.gif')
 
-    return zip(imgs_switched, masks)
+    return zip(imgs_normalized, masks)
diff --git a/myloss.py b/myloss.py
index e65f0f1..2c26c39 100644
--- a/myloss.py
+++ b/myloss.py
@@ -1,34 +1,52 @@
+
+#
+# myloss.py : implementation of the Dice coeff and the associated loss
+#
+
 import torch
-from torch.nn.modules.loss import _Loss
-from torch.autograd import Function
 import torch.nn.functional as F
 
+from torch.nn.modules.loss import _Loss
+from torch.autograd import Function, Variable
+
+
 class DiceCoeff(Function):
+    """Dice coeff for individual examples"""
+    def forward(self, input, target):
+        self.save_for_backward(input, target)
+        self.inter = torch.dot(input, target) + 0.0001
+        self.union = torch.sum(input) + torch.sum(target) + 0.0001
 
-    def forward(ctx, input, target):
-        ctx.save_for_backward(input, target)
-        ctx.inter = torch.dot(input, target) + 0.0001
-        ctx.union = torch.sum(input) + torch.sum(target) + 0.0001
-
-        t = 2*ctx.inter.float()/ctx.union.float()
+        t = 2*self.inter.float()/self.union.float()
         return t
 
     # This function has only a single output, so it gets only one gradient
-    def backward(ctx, grad_output):
+    def backward(self, grad_output):
 
-        input, target = ctx.saved_variables
+        input, target = self.saved_variables
         grad_input = grad_target = None
 
         if self.needs_input_grad[0]:
-            grad_input = grad_output * 2 * (target * ctx.union + ctx.inter) \
-                         / ctx.union * ctx.union
+            grad_input = grad_output * 2 * (target * self.union + self.inter) \
+                         / self.union * self.union
         if self.needs_input_grad[1]:
             grad_target = None
 
         return grad_input, grad_target
 
+
 def dice_coeff(input, target):
-    return DiceCoeff().forward(input, target)
+    """Dice coeff for batches"""
+    if input.is_cuda:
+        s = Variable(torch.FloatTensor(1).cuda().zero_())
+    else:
+        s = Variable(torch.FloatTensor(1).zero_())
+
+    for i, c in enumerate(zip(input, target)):
+        s = s + DiceCoeff().forward(c[0], c[1])
+
+    return s / (i+1)
+
 
 class DiceLoss(_Loss):
     def forward(self, input, target):
diff --git a/train.py b/train.py
index e69de29..2dedee2 100644
--- a/train.py
+++ b/train.py
@@ -0,0 +1,105 @@
+import torch
+
+from load import *
+from data_vis import *
+from utils import split_train_val, batch
+from myloss import DiceLoss
+from unet_model import UNet
+from torch.autograd import Variable
+from torch import optim
+from optparse import OptionParser
+
+
+def train_net(net, epochs=5, batch_size=2, lr=0.1, val_percent=0.05,
+              cp=True, gpu=False):
+    dir_img = 'data/train/'
+    dir_mask = 'data/train_masks/'
+    dir_checkpoint = 'checkpoints/'
+
+    # get ids
+    ids = get_ids(dir_img)
+    ids = split_ids(ids)
+
+    iddataset = split_train_val(ids, val_percent)
+
+    print('''
+    Starting training:
+        Epochs: {}
+        Batch size: {}
+        Learning rate: {}
+        Training size: {}
+        Validation size: {}
+        Checkpoints: {}
+        CUDA: {}
+    '''.format(epochs, batch_size, lr, len(iddataset['train']),
+               len(iddataset['val']), str(cp), str(gpu)))
+
+    N_train = len(iddataset['train'])
+
+    train = get_imgs_and_masks(iddataset['train'], dir_img, dir_mask)
+    val = get_imgs_and_masks(iddataset['val'], dir_img, dir_mask)
+
+    optimizer = optim.Adam(net.parameters(), lr=lr)
+    criterion = DiceLoss()
+
+    for epoch in range(epochs):
+        print('Starting epoch {}/{}.'.format(epoch+1, epochs))
+
+        epoch_loss = 0
+
+        for i, b in enumerate(batch(train, batch_size)):
+            X = np.array([i[0] for i in b])
+            y = np.array([i[1] for i in b])
+
+            X = torch.FloatTensor(X)
+            y = torch.ByteTensor(y)
+
+            if gpu:
+                X = Variable(X).cuda()
+                y = Variable(y).cuda()
+            else:
+                X = Variable(X)
+                y = Variable(y)
+
+            optimizer.zero_grad()
+
+            y_pred = net(X)
+
+            loss = criterion(y_pred, y.float())
+            epoch_loss += loss.data[0]
+
+            print('{0:.4f} --- loss: {1:.6f}'.format(i*batch_size/N_train,
+                                              loss.data[0]))
+
+            loss.backward()
+            optimizer.step()
+
+        print('Epoch finished ! Loss: {}'.format(epoch_loss/i))
+
+        if cp:
+            torch.save(net.state_dict(),
+                       dir_checkpoint + 'CP{}.pth'.format(epoch+1))
+
+            print('Checkpoint {} saved !'.format(epoch+1))
+
+
+parser = OptionParser()
+parser.add_option("-e", "--epochs", dest="epochs", default=5, type="int",
+                  help="number of epochs")
+parser.add_option("-b", "--batch-size", dest="batchsize", default=10,
+                  type="int", help="batch size")
+parser.add_option("-l", "--learning-rate", dest="lr", default=0.1,
+                  type="int", help="learning rate")
+parser.add_option("-g", "--gpu", action="store_true", dest="gpu",
+                  default=False, help="use cuda")
+parser.add_option("-n", "--ngpu", action="store_false", dest="gpu",
+                  default=False, help="use cuda")
+
+
+(options, args) = parser.parse_args()
+
+net = UNet(3, 1)
+if options.gpu:
+    net.cuda()
+
+train_net(net, options.epochs, options.batchsize, options.lr, gpu=options.gpu)
diff --git a/unet_model.py b/unet_model.py
index 6d5dc39..5129e64 100644
--- a/unet_model.py
+++ b/unet_model.py
@@ -4,6 +4,7 @@ import torch.nn.functional as F
 
 from unet_parts import *
 
+
 class UNet(nn.Module):
     def __init__(self, n_channels, n_classes):
         super(UNet, self).__init__()
diff --git a/unet_parts.py b/unet_parts.py
index 37ec10e..08fd8d0 100644
--- a/unet_parts.py
+++ b/unet_parts.py
@@ -4,6 +4,7 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
+
 class double_conv(nn.Module):
     def __init__(self, in_ch, out_ch):
         super(double_conv, self).__init__()
@@ -13,10 +14,12 @@ class double_conv(nn.Module):
             nn.Conv2d(out_ch, out_ch, 3, padding=1),
             nn.ReLU()
         )
+
     def forward(self, x):
         x = self.conv(x)
         return x
 
+
 class inconv(nn.Module):
     def __init__(self, in_ch, out_ch):
         super(inconv, self).__init__()
@@ -26,6 +29,7 @@ class inconv(nn.Module):
         x = self.conv(x)
         return x
 
+
 class down(nn.Module):
     def __init__(self, in_ch, out_ch):
         super(down, self).__init__()
@@ -38,15 +42,15 @@ class down(nn.Module):
         x = self.mpconv(x)
         return x
 
+
 class up(nn.Module):
     def __init__(self, in_ch, out_ch):
         super(up, self).__init__()
         self.up = nn.UpsamplingBilinear2d(scale_factor=2)
-        #self.up = nn.ConvTranspose2d(in_ch, out_ch, 2, stride=2)
+        # self.up = nn.ConvTranspose2d(in_ch, out_ch, 2, stride=2)
         self.conv = double_conv(in_ch, out_ch)
 
     def forward(self, x1, x2):
-
         x1 = self.up(x1)
         diffX = x1.size()[2] - x2.size()[2]
         diffY = x1.size()[3] - x2.size()[3]
@@ -56,6 +60,7 @@ class up(nn.Module):
         x = self.conv(x)
         return x
 
+
 class outconv(nn.Module):
     def __init__(self, in_ch, out_ch):
         super(outconv, self).__init__()
diff --git a/utils.py b/utils.py
index 9732526..b51cd74 100644
--- a/utils.py
+++ b/utils.py
@@ -1,26 +1,54 @@
 import PIL
 import numpy as np
+import random
+
 
 def get_square(img, pos):
-    """Extract a left or a right square from PILimg"""
-    """shape : (H, W, C))"""
+    """Extract a left or a right square from PILimg shape : (H, W, C))"""
     img = np.array(img)
-
     h = img.shape[0]
-    w = img.shape[1]
-
     if pos == 0:
         return img[:, :h]
     else:
         return img[:, -h:]
 
-def resize_and_crop(pilimg, scale=0.5, final_height=640):
+
+def resize_and_crop(pilimg, scale=0.2, final_height=None):
     w = pilimg.size[0]
     h = pilimg.size[1]
     newW = int(w * scale)
     newH = int(h * scale)
-    diff = newH - final_height
+
+    if not final_height:
+        diff = 0
+    else:
+        diff = newH - final_height
 
     img = pilimg.resize((newW, newH))
     img = img.crop((0, diff // 2, newW, newH - diff // 2))
     return img
+
+
+def batch(iterable, batch_size):
+    """Yields lists by batch"""
+    b = []
+    for i, t in enumerate(iterable):
+        b.append(t)
+        if (i+1) % batch_size == 0:
+            yield b
+            b = []
+
+    if len(b) > 0:
+        yield b
+
+
+def split_train_val(dataset, val_percent=0.05):
+    dataset = list(dataset)
+    length = len(dataset)
+    n = int(length * val_percent)
+    random.shuffle(dataset)
+    return {'train': dataset[:-n], 'val': dataset[-n:]}
+
+
+def normalize(x):
+    return x / 255