test.py 3.92 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34
#!/usr/bin/env python
# encoding: utf-8


""" Unit tests

"""

import numpy
import torch

def test_architectures():

  a = numpy.random.rand(1, 3, 128, 128).astype("float32")
  t = torch.from_numpy(a)
  
  number_of_classes = 20
  output_dimension = number_of_classes
  
  # CASIANet
  from ..architectures import CASIANet
  net = CASIANet(number_of_classes)
  embedding_dimension = 320
  output, emdedding = net.forward(t)
  assert output.shape == torch.Size([1, 20])
  assert emdedding.shape == torch.Size([1, 320])
  
  # CNN8
  from ..architectures import CNN8
  net = CNN8(number_of_classes)
  embedding_dimension = 512
  output, emdedding = net.forward(t)
  assert output.shape == torch.Size([1, 20])
  assert emdedding.shape == torch.Size([1, 512])
35 36 37 38 39 40 41 42 43 44 45 46 47 48 49

  # DCGAN
  d = numpy.random.rand(1, 3, 64, 64).astype("float32")
  t = torch.from_numpy(d)
  from ..architectures import DCGAN_discriminator
  discriminator = DCGAN_discriminator(1)
  output = discriminator.forward(t)
  assert output.shape == torch.Size([1])

  g = numpy.random.rand(1, 100, 1, 1).astype("float32")
  t = torch.from_numpy(g)
  from ..architectures import DCGAN_generator
  generator = DCGAN_generator(1)
  output = generator.forward(t)
  assert output.shape == torch.Size([1, 3, 64, 64])
50 51 52 53
 

def test_transforms():

54
  image = numpy.random.rand(3, 128, 128).astype("uint8")
55 56

  from ..datasets import RollChannels 
57
  sample = {'image': image}
58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102
  rc = RollChannels()
  rc(sample)
  assert sample['image'].shape == (128, 128, 3)

  from ..datasets import ToTensor 
  tt = ToTensor()
  tt(sample)
  assert isinstance(sample['image'], torch.Tensor)

  from ..datasets import Normalize 
  image_copy = torch.Tensor(sample['image'])
  norm = Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
  norm(sample)
  for c in range(3):
    for h in range(sample['image'].shape[0]):
      for w in range(sample['image'].shape[1]):
        assert (abs(sample['image'][c, h, w]) - abs((image_copy[c, h, w] - 0.5) / 0.5)) < 1e-10


def test_map_labels():

  labels = ['1', '4', '7']
  from ..datasets import map_labels
  new_labels = map_labels(labels)
  new_labels = sorted(new_labels)
  assert new_labels == ['0', '1', '2']

  new_labels = map_labels(labels, start_index = 5)
  new_labels = sorted(new_labels)
  assert new_labels == ['5', '6', '7']
 

from torch.utils.data import Dataset
class DummyDataSet(Dataset):
  def __init__(self):
    pass
  def __len__(self):
    return 100
  def __getitem__(self, idx):
    data =  numpy.random.rand(3, 128, 128).astype("float32")
    label = numpy.random.randint(20)
    sample = {'image': torch.from_numpy(data), 'label': label}
    return sample 


103
def test_CNNtrainer():
104 105 106 107 108 109 110 111 112 113 114 115 116 117

  from ..architectures import CNN8
  net = CNN8(20)

  dataloader = torch.utils.data.DataLoader(DummyDataSet(), batch_size=32, shuffle=True)
  
  from ..trainers import CNNTrainer
  trainer = CNNTrainer(net, verbosity_level=3)
  trainer.train(dataloader, n_epochs=1, output_dir='.')

  import os
  assert os.path.isfile('model_1_0.pth')

  os.remove('model_1_0.pth')
118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151


class DummyDataSetGAN(Dataset):
  def __init__(self):
    pass
  def __len__(self):
    return 100
  def __getitem__(self, idx):
    data =  numpy.random.rand(3, 64, 64).astype("float32")
    sample = {'image': torch.from_numpy(data)}
    return sample 

def test_DCGANtrainer():

  from ..architectures import DCGAN_generator
  from ..architectures import DCGAN_discriminator
  g = DCGAN_generator(1)
  d = DCGAN_discriminator(1)

  dataloader = torch.utils.data.DataLoader(DummyDataSetGAN(), batch_size=32, shuffle=True)
  
  from ..trainers import DCGANTrainer
  trainer = DCGANTrainer(g, d, batch_size=32, noise_dim=100, use_gpu=False, verbosity_level=2)
  trainer.train(dataloader, n_epochs=1, output_dir='.')

  import os
  assert os.path.isfile('fake_samples_epoch_000.png')
  assert os.path.isfile('netD_epoch_0.pth')
  assert os.path.isfile('netG_epoch_0.pth')

  os.remove('fake_samples_epoch_000.png')
  os.remove('netD_epoch_0.pth')
  os.remove('netG_epoch_0.pth')