def test_proc(self) -> None:
test_a_paths = file_util.get_file_paths(str(self.test_data_root) +
str(self.test_a_path),
verbose=self.verbose)
assert len(test_a_paths) != 0
test_b_paths = file_util.get_file_paths(str(self.test_data_root) +
str(self.test_b_path),
verbose=self.verbose)
assert len(test_b_paths) != 0
# For this test we need an equal amount of data
dataset_size = min(len(test_a_paths), len(test_b_paths))
test_a_paths = test_a_paths[0:dataset_size]
test_b_paths = test_b_paths[0:dataset_size]
print('Processing %d A images, %d B images from paths [%s]. [%s]' %\
(len(test_a_paths), len(test_b_paths),
str(self.test_data_root) + str(self.test_a_path),
str(self.test_data_root) + str(self.test_b_path))
)
test_image_dataset_name = 'images'
align_proc = aligned_data_proc.AlignedImageJoin(
image_dataset_name=test_image_dataset_name,
image_shape=self.test_image_shape,
verbose=self.verbose)
test_outfile = 'data/test_aligned_image_join.h5'
align_proc.proc(test_a_paths, test_b_paths, test_outfile)
# check the dataset contents
with h5py.File(test_outfile, 'r') as fp:
dataset_keys = fp.keys()
assert test_image_dataset_name in dataset_keys
assert 'a_ids' in dataset_keys
assert 'b_ids' in dataset_keys
## Test as HDF5Dataset
# Note that I am being sneaky here as I am loading the image as both
# the feature and the label. This is only done to shut the system up,
# since we don't actually need a label here. For the purpose of this
# test we care that the output is in fact twice the width (because that
# means that we joined the images correctly).
test_dataset = hdf5_dataset.HDF5Dataset(
test_outfile,
feature_name=align_proc.image_dataset_name,
label_name=align_proc.image_dataset_name,
verbose=self.verbose)
assert len(test_a_paths) == len(test_dataset)
# Check that the elements in the dataset are the expected size
for elem_idx, (feature, _) in enumerate(test_dataset):
print('Checking element [%d / %d]' %
(elem_idx + 1, len(test_dataset)),
end='\r')
assert self.test_out_image_shape == feature.shape
print('\n OK')
if self.remove:
os.remove(test_outfile)
def get_dataset(image_size: int = 64):
dataset = hdf5_dataset.HDF5Dataset(
DATASET_ROOT,
feature_name='images',
label_name='labels',
#transform = gan_data_transform
)
return dataset
def get_hdf5_dataset(dataset_path: str,
feature_name: str = 'images',
label_name: str = 'labels') -> hdf5_dataset.HDF5Dataset:
train_dataset = hdf5_dataset.HDF5Dataset(
GLOBAL_OPTS['dataset'],
feature_name='images',
label_name='labels',
)
return train_dataset
def translate_dataset(max_images: int = 128) -> None:
inferrer = pix2pix_inferrer.Pix2PixInferrer(
device_id=GLOBAL_OPTS['device_id'])
inferrer.load_model(GLOBAL_OPTS['checkpoint'])
dataset = hdf5_dataset.HDF5Dataset(GLOBAL_OPTS['dataset_path'],
feature_name='images',
label_name='labels')
data_loader = torch.utils.data.DataLoader(dataset,
batch_size=1,
shuffle=False)
# NOTE : for now, we just use the batch index as part of the filename
fig, ax = plt.subplots() # Is it faster to move this out of the loop?
img_idx = 0
for batch_idx, (data_a, data_b) in enumerate(data_loader):
print('Processing image [%d / %d]' %
(batch_idx + 1, data_loader.batch_size * len(data_loader)),
end='\r')
fake_a = inferrer.forward(data_a)
if data_loader.batch_size > 1:
for n in range(fake_a.shape[0]):
fname = os.path.splitext(
GLOBAL_OPTS['outfile'])[0] + '_' + str(img_idx) + '.png'
write_img_tensor(fig, ax, fname, fake_a[n])
img_idx += 1
else:
fname = os.path.splitext(
GLOBAL_OPTS['outfile'])[0] + '_' + str(img_idx) + '.png'
write_img_tensor(fig, ax, fname, fake_a)
img_idx += 1
if img_idx >= max_images:
break
print('\n done')
def main() -> None:
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
train_dataset_transform = transforms.Compose([
transforms.RandomRotation(5),
transforms.RandomHorizontalFlip(),
transforms.RandomResizedCrop(224,
scale=(0.96, 1.0),
ratio=(0.95, 1.05)),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
test_dataset_transform = transforms.Compose([
transforms.Resize([224, 224]),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
# HDF5 Datasets
if GLOBAL_USE_HDF5 is True:
cvd_train_dataset = hdf5_dataset.HDF5Dataset(
GLOBAL_OPTS['train_dataset'],
feature_name='images',
label_name='labels',
label_max_dim=1,
transform=normalize)
cvd_val_dataset = hdf5_dataset.HDF5Dataset(GLOBAL_OPTS['test_dataset'],
feature_name='images',
label_name='labels',
label_max_dim=1,
transform=normalize)
else:
cvd_train_dir = '/home/kreshnik/ml-data/cats-vs-dogs/train'
# ImageFolder dataset
cvd_train_dataset = datasets.ImageFolder(cvd_train_dir,
train_dataset_transform)
csv_val_dir = '/home/kreshnik/ml-data/cats-vs-dogs/test'
cvd_val_dataset = datasets.ImageFolder(csv_val_dir,
test_dataset_transform)
# get a network
model = cvdnet.CVDNet2()
cvd_train = cvd_trainer.CVDTrainer(
model,
# dataset options
train_dataset=cvd_train_dataset,
val_dataset=cvd_val_dataset,
# training options
loss_function='CrossEntropyLoss',
learning_rate=GLOBAL_OPTS['learning_rate'],
weight_decay=GLOBAL_OPTS['weight_decay'],
momentum=GLOBAL_OPTS['momentum'],
num_epochs=GLOBAL_OPTS['num_epochs'],
batch_size=GLOBAL_OPTS['batch_size'],
val_batch_size=GLOBAL_OPTS['val_batch_size'],
# checkpoint
checkpoint_name=GLOBAL_OPTS['checkpoint_name'],
save_every=GLOBAL_OPTS['save_every'],
# device
device_id=GLOBAL_OPTS['device_id'],
# other
print_every=GLOBAL_OPTS['print_every'],
verbose=GLOBAL_OPTS['verbose'])
# Add a tensorboard writer
if GLOBAL_OPTS['tensorboard_dir'] is not None:
if not os.path.isdir(GLOBAL_OPTS['tensorboard_dir']):
os.mkdir(GLOBAL_OPTS['tensorboard_dir'])
writer = tensorboard.SummaryWriter(
log_dir=GLOBAL_OPTS['tensorboard_dir'])
cvd_train.set_tb_writer(writer)
# Optionally do a search pass here and add a scheduler
if GLOBAL_OPTS['find_lr']:
lr_finder = expr_util.get_lr_finder(cvd_train)
lr_find_start_time = time.time()
lr_finder.find()
lr_find_min, lr_find_max = lr_finder.get_lr_range()
lr_find_end_time = time.time()
lr_find_total_time = lr_find_end_time - lr_find_start_time
print('Found learning rate range %.4f -> %.4f' %
(lr_find_min, lr_find_max))
print('Total find time [%s] ' %\
str(timedelta(seconds = lr_find_total_time))
)
# Now get a scheduler
stepsize = cvd_train.get_num_epochs() * len(trainer.train_loader) // 2
# get scheduler
lr_scheduler = expr_util.get_scheduler(
lr_find_min,
lr_find_max,
stepsize,
sched_type='TriangularScheduler')
cvd_train.set_lr_scheduler(lr_scheduler)
# train the model
train_start_time = time.time()
cvd_train.train()
train_end_time = time.time()
train_total_time = train_end_time - train_start_time
print('Total scheduled training time [%s] (%d epochs) %s' %\
(repr(cvd_train), cvd_train.cur_epoch,
str(timedelta(seconds = train_total_time)))
)
# Show results
fig, ax = vis_loss_history.get_figure_subplots(num_subplots=2)
vis_loss_history.plot_train_history_2subplots(
ax,
cvd_train.get_loss_history(),
acc_history=cvd_train.get_acc_history(),
iter_per_epoch=cvd_train.iter_per_epoch,
loss_title='CVD Loss',
acc_title='CVD Acc',
cur_epoch=cvd_train.cur_epoch)
fig.savefig('figures/cvd_train.png')
def main() -> None:
if GLOBAL_OPTS['dataset'] is not None:
train_dataset = hdf5_dataset.HDF5Dataset(
GLOBAL_OPTS['dataset'],
feature_name='images',
label_name='labels',
)
else:
gan_data_transform = transforms.Compose([
transforms.Resize(GLOBAL_OPTS['image_size']),
transforms.CenterCrop(GLOBAL_OPTS['image_size']),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
])
train_dataset = datasets.ImageFolder(root=GLOBAL_OPTS['dataset_root'],
transform=gan_data_transform)
# get some models
generator = dcgan.DCGANGenerator(zvec_dim=GLOBAL_OPTS['zvec_dim'],
num_filters=GLOBAL_OPTS['g_num_filters'],
img_size=GLOBAL_OPTS['image_size'])
discriminator = dcgan.DCGANDiscriminator(
num_filters=GLOBAL_OPTS['d_num_filters'],
img_size=GLOBAL_OPTS['image_size'])
# get a trainer
gan_trainer = dcgan_trainer.DCGANTrainer(
discriminator,
generator,
# DCGAN trainer specific arguments
beta1=GLOBAL_OPTS['beta1'],
# general trainer arguments
train_dataset=train_dataset,
# training opts
learning_rate=GLOBAL_OPTS['learning_rate'],
num_epochs=GLOBAL_OPTS['num_epochs'],
batch_size=GLOBAL_OPTS['batch_size'],
# Checkpoints
save_every=GLOBAL_OPTS['save_every'],
checkpoint_name=GLOBAL_OPTS['checkpoint_name'],
checkpoint_dir=GLOBAL_OPTS['checkpoint_dir'],
# display
print_every=GLOBAL_OPTS['print_every'],
verbose=GLOBAL_OPTS['verbose'],
# device
device_id=GLOBAL_OPTS['device_id'])
# add a summary writer
if GLOBAL_OPTS['tensorboard_dir'] is not None:
writer = tensorboard.SummaryWriter(
log_dir=GLOBAL_OPTS['tensorboard_dir'])
gan_trainer.set_tb_writer(writer)
# load a checkpoint?
if GLOBAL_OPTS['load_checkpoint'] is not None:
gan_trainer.load_checkpoint(GLOBAL_OPTS['load_checkpoint'])
print(gan_trainer.device)
# Get a scheduler
lr_scheduler = schedule.DecayToEpoch(
non_decay_time=int(GLOBAL_OPTS['num_epochs'] // 2),
decay_length=int(GLOBAL_OPTS['num_epochs'] // 2),
initial_lr=GLOBAL_OPTS['learning_rate'],
final_lr=0.0)
print('Created scheduler')
print(lr_scheduler)
gan_trainer.set_lr_scheduler(lr_scheduler)
train_start_time = time.time()
gan_trainer.train()
train_end_time = time.time()
train_total_time = train_end_time - train_start_time
print('Total training time : %s' %
str(timedelta(seconds=train_total_time)))
# show the training results
dcgan_fig, dcgan_ax = vis_loss_history.get_figure_subplots(1)
vis_loss_history.plot_train_history_dcgan(
dcgan_ax,
gan_trainer.get_g_loss_history(),
gan_trainer.get_d_loss_history(),
cur_epoch=gan_trainer.cur_epoch,
iter_per_epoch=gan_trainer.iter_per_epoch)
dcgan_fig.tight_layout()
dcgan_fig.savefig('figures/dcgan_train_history.png')