| @@ -0,0 +1,132 @@ | |||||
| import time | |||||
| import multiprocessing | |||||
| import tensorflow as tf | |||||
| import os | |||||
| os.environ['TL_BACKEND'] = 'tensorflow' | |||||
| import tensorlayer as tl | |||||
| from .densenet import densenet | |||||
| tl.logging.set_verbosity(tl.logging.DEBUG) | |||||
| X_train, y_train, X_test, y_test = tl.files.load_cifar10_dataset(shape=(-1, 32, 32, 3), plotable=False) | |||||
| # get the network | |||||
| net = densenet("densenet-100") | |||||
| # training settings | |||||
| batch_size = 128 | |||||
| n_epoch = 500 | |||||
| learning_rate = 0.0001 | |||||
| print_freq = 5 | |||||
| n_step_epoch = int(len(y_train) / batch_size) | |||||
| n_step = n_epoch * n_step_epoch | |||||
| shuffle_buffer_size = 128 | |||||
| train_weights = net.trainable_weights | |||||
| optimizer = tl.optimizers.Adam(learning_rate) | |||||
| metrics = tl.metric.Accuracy() | |||||
| def generator_train(): | |||||
| inputs = X_train | |||||
| targets = y_train | |||||
| if len(inputs) != len(targets): | |||||
| raise AssertionError("The length of inputs and targets should be equal") | |||||
| for _input, _target in zip(inputs, targets): | |||||
| # yield _input.encode('utf-8'), _target.encode('utf-8') | |||||
| yield _input, _target | |||||
| def generator_test(): | |||||
| inputs = X_test | |||||
| targets = y_test | |||||
| if len(inputs) != len(targets): | |||||
| raise AssertionError("The length of inputs and targets should be equal") | |||||
| for _input, _target in zip(inputs, targets): | |||||
| # yield _input.encode('utf-8'), _target.encode('utf-8') | |||||
| yield _input, _target | |||||
| def _map_fn_train(img, target): | |||||
| # 1. Randomly crop a [height, width] section of the image. | |||||
| img = tf.image.random_crop(img, [24, 24, 3]) | |||||
| # 2. Randomly flip the image horizontally. | |||||
| img = tf.image.random_flip_left_right(img) | |||||
| # 3. Randomly change brightness. | |||||
| img = tf.image.random_brightness(img, max_delta=63) | |||||
| # 4. Randomly change contrast. | |||||
| img = tf.image.random_contrast(img, lower=0.2, upper=1.8) | |||||
| # 5. Subtract off the mean and divide by the variance of the pixels. | |||||
| img = tf.image.per_image_standardization(img) | |||||
| target = tf.reshape(target, ()) | |||||
| return img, target | |||||
| def _map_fn_test(img, target): | |||||
| # 1. Crop the central [height, width] of the image. | |||||
| img = tf.image.resize_with_pad(img, 24, 24) | |||||
| # 2. Subtract off the mean and divide by the variance of the pixels. | |||||
| img = tf.image.per_image_standardization(img) | |||||
| img = tf.reshape(img, (24, 24, 3)) | |||||
| target = tf.reshape(target, ()) | |||||
| return img, target | |||||
| # dataset API and augmentation | |||||
| train_ds = tf.data.Dataset.from_generator( | |||||
| generator_train, output_types=(tf.float32, tf.int32) | |||||
| ) # , output_shapes=((24, 24, 3), (1))) | |||||
| train_ds = train_ds.map(_map_fn_train,num_parallel_calls=multiprocessing.cpu_count()) | |||||
| # train_ds = train_ds.repeat(n_epoch) | |||||
| train_ds = train_ds.shuffle(shuffle_buffer_size) | |||||
| train_ds = train_ds.prefetch(buffer_size=4096) | |||||
| train_ds = train_ds.batch(batch_size) | |||||
| # value = train_ds.make_one_shot_iterator().get_next() | |||||
| test_ds = tf.data.Dataset.from_generator( | |||||
| generator_test, output_types=(tf.float32, tf.int32) | |||||
| ) # , output_shapes=((24, 24, 3), (1))) | |||||
| # test_ds = test_ds.shuffle(shuffle_buffer_size) | |||||
| test_ds = test_ds.map(_map_fn_test,num_parallel_calls=multiprocessing.cpu_count()) | |||||
| # test_ds = test_ds.repeat(n_epoch) | |||||
| test_ds = test_ds.prefetch(buffer_size=4096) | |||||
| test_ds = test_ds.batch(batch_size) | |||||
| # value_test = test_ds.make_one_shot_iterator().get_next() | |||||
| class WithLoss(tl.layers.Module): | |||||
| def __init__(self, net, loss_fn): | |||||
| super(WithLoss, self).__init__() | |||||
| self._net = net | |||||
| self._loss_fn = loss_fn | |||||
| def forward(self, data, label): | |||||
| out = self._net(data) | |||||
| loss = self._loss_fn(out, label) | |||||
| return loss | |||||
| net_with_loss = WithLoss(net, loss_fn=tl.cost.softmax_cross_entropy_with_logits) | |||||
| net_with_train = tl.models.TrainOneStep(net_with_loss, optimizer, train_weights) | |||||
| for epoch in range(n_epoch): | |||||
| start_time = time.time() | |||||
| net.set_train() | |||||
| train_loss, train_acc, n_iter = 0, 0, 0 | |||||
| for X_batch, y_batch in train_ds: | |||||
| X_batch = tl.ops.convert_to_tensor(X_batch.numpy(), dtype=tl.float32) | |||||
| y_batch = tl.ops.convert_to_tensor(y_batch.numpy(), dtype=tl.int64) | |||||
| _loss_ce = net_with_train(X_batch, y_batch) | |||||
| train_loss += _loss_ce | |||||
| n_iter += 1 | |||||
| _logits = net(X_batch) | |||||
| metrics.update(_logits, y_batch) | |||||
| train_acc += metrics.result() | |||||
| metrics.reset() | |||||
| print("Epoch {} of {} took {}".format(epoch + 1, n_epoch, time.time() - start_time)) | |||||
| print(" train loss: {}".format(train_loss / n_iter)) | |||||
| print(" train acc: {}".format(train_acc / n_iter)) | |||||