[script] first version a the full train script, achieving some results

bob/learn/pytorch/scripts/train.py

@@ -13,7 +13,7 @@ Options:
   -h, --help                Show this screen.
   -V, --version             Show version.
   -l, --latent-dim=<int>    the dimension of the encoded ID [default: 320]
-  -b, --batch-size=<int>    The size of your mini-batch [default: 128]
+  -b, --batch-size=<int>    The size of your mini-batch [default: 64]
   -e, --epochs=<int>        The number of training epochs [default: 100] 
   -s, --sample=<int>        Save generated images at every 'sample' batch iteration [default: 100000000000] 
   -o, --output-dir=<path>   Dir to save the logs, models and images [default: ./drgan-light-mpie-casia/] 
@@ -40,9 +40,21 @@ from docopt import docopt
 version = pkg_resources.require('bob.learn.pytorch')[0].version
 
 import numpy
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import torchvision.transforms as transforms
+import torchvision.utils as vutils
+
+
+from torch.autograd import Variable
+
 
 from bob.learn.pytorch.datasets.multipie import MultiPIEDataset
-#import bob.learn.pytorch.datasets.multipie
+from bob.learn.pytorch.datasets.multipie import RollChannels
+from bob.learn.pytorch.architectures.DCGAN import _netG
+from bob.learn.pytorch.architectures.DCGAN import _netD
+from bob.learn.pytorch.architectures.DCGAN import weights_init
 
 def main(user_input=None):
   
@@ -70,15 +82,119 @@ def main(user_input=None):
   log_dir = os.path.join(output_dir, 'logs')
   model_dir = os.path.join(output_dir, 'models')
 
-  #face_dataset = MultiPIEDataset(root_dir='/idiap/resource/database/Multi-Pie/data/')
-  face_dataset = MultiPIEDataset(root_dir='/idiap/temp/heusch/data/multipie-cropped-64x64')
-  #print len(face_dataset)
-
-  from matplotlib import pyplot
-  for i in range(len(face_dataset)):
-    sample = face_dataset[i]
-    pyplot.title('Sample {}: ID -> {}, pose ->{}'.format(i, sample['id'], sample['pose']))
-    pyplot.imshow(numpy.rollaxis(numpy.rollaxis(sample['image'], 2),2))
-    pyplot.show()
-    
+  try:
+    os.makedirs(output_dir)
+  except OSError:
+    pass
+
+
+  # data 
+  # WARNING with the transforms ... act on labels too, at some point, I may have to write my own
+  # Also, in 'ToTensor', there is a reshape performed from: HxWxC to CxHxW
+  face_dataset = MultiPIEDataset(root_dir='/idiap/temp/heusch/data/multipie-cropped-64x64', frontal_only=True, 
+      transform=transforms.Compose([RollChannels(), transforms.ToTensor(),transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]))
+  logger.info("There are {} training images".format(len(face_dataset))) 
+
+  dataloader = torch.utils.data.DataLoader(face_dataset, batch_size=batch_size, shuffle=True)
+
+
+  #from matplotlib import pyplot
+  #for i in range(len(face_dataset)):
+  #  sample = face_dataset[i]
+  #  pyplot.title('Sample {}: ID -> {}, pose ->{}'.format(i, sample['id'], sample['pose']))
+  #  pyplot.imshow(sample['image'])
+  #  #pyplot.imshow(numpy.rollaxis(numpy.rollaxis(sample['image'], 2),2))
+  #  pyplot.show()
+   
+  # network
+  ngpu = 1
+  netG = _netG(ngpu)
+  netG.apply(weights_init)
+  print(netG)
+ 
+  netD = _netD(ngpu)
+  netD.apply(weights_init)
+  print(netD)
+
+  criterion = nn.BCELoss()
+
+  nz = 100
+  input = torch.FloatTensor(batch_size, 3, 64, 64)
+  noise = torch.FloatTensor(batch_size, nz, 1, 1)
+  fixed_noise = torch.FloatTensor(batch_size, nz, 1, 1).normal_(0, 1)
+  label = torch.FloatTensor(batch_size)
+  real_label = 1
+  fake_label = 0
+
+  fixed_noise = Variable(fixed_noise)
+
+  # setup optimizer
+  lr = 0.0002
+  beta1 = 0.5
+  optimizerD = optim.Adam(netD.parameters(), lr=lr, betas=(beta1, 0.999))
+  optimizerG = optim.Adam(netG.parameters(), lr=lr, betas=(beta1, 0.999))
+
+  niter = 10
+  for epoch in range(niter):
+    for i, data in enumerate(dataloader, 0):
+     
+      #print data
+      #print len(data)
+
+      ############################
+      # (1) Update D network: maximize log(D(x)) + log(1 - D(G(z)))
+      ###########################
+      # train with real
+      netD.zero_grad()
+      real_cpu = data
+      #print (type(real_cpu))
+      #print (real_cpu)
+
+      batch_size = real_cpu.size(0)
+      input.resize_as_(real_cpu).copy_(real_cpu)
+      label.resize_(batch_size).fill_(real_label)
+      inputv = Variable(input)
+      labelv = Variable(label)
+
+      output = netD(inputv)
+      errD_real = criterion(output, labelv)
+      errD_real.backward()
+      D_x = output.data.mean()
+
+      # train with fake
+      noise.resize_(batch_size, nz, 1, 1).normal_(0, 1)
+      noisev = Variable(noise)
+      fake = netG(noisev)
+      labelv = Variable(label.fill_(fake_label))
+      output = netD(fake.detach())
+      errD_fake = criterion(output, labelv)
+      errD_fake.backward()
+      D_G_z1 = output.data.mean()
+      errD = errD_real + errD_fake
+      optimizerD.step()
+
+      ############################
+      # (2) Update G network: maximize log(D(G(z)))
+      ###########################
+      netG.zero_grad()
+      labelv = Variable(label.fill_(real_label))  # fake labels are real for generator cost
+      output = netD(fake)
+      errG = criterion(output, labelv)
+      errG.backward()
+      D_G_z2 = output.data.mean()
+      optimizerG.step()
+
+      print'[%d/%d][%d/%d] Loss_D: %.4f Loss_G: %.4f D(x): %.4f D(G(z)): %.4f / %.4f' % (epoch, niter, i, len(dataloader), errD.data[0], errG.data[0], D_x, D_G_z1, D_G_z2)
+      if i % 100 == 0:
+        vutils.save_image(real_cpu,
+                '%s/real_samples.png' % output_dir,
+                normalize=True)
+        fake = netG(fixed_noise)
+        vutils.save_image(fake.data,
+                '%s/fake_samples_epoch_%03d.png' % (output_dir, epoch),
+                normalize=True)
+
+    # do checkpointing
+    torch.save(netG.state_dict(), '%s/netG_epoch_%d.pth' % (output_dir, epoch))
+    torch.save(netD.state_dict(), '%s/netD_epoch_%d.pth' % (output_dir, epoch))