add cifar and lincls

Author: sungchul2

.gitignore

@@ -0,0 +1,5 @@
+.ipynb_checkpoints/
+__pycache__/
+result/
+*.ipynb
+*.sh

augment.py

@@ -18,14 +18,29 @@ def _augment_simsiam(self, x, shape, coord=[[[0., 0., 1., 1.]]]):
         x = self._resize(x)
         x = self._random_color_jitter(x, p=.8)
         x = self._random_grayscale(x, p=.2)
-        x = self._random_gaussian_blur(x, p=.5)
+        if self.args.dataset == 'imagenet':
+            x = self._random_gaussian_blur(x, p=.5)
         x = self._random_hflip(x)
         x = self._standardize(x)
         return x
 
     def _augment_lincls(self, x, shape, coord=[[[0., 0., 1., 1.]]]):
-        x = self._crop(x, shape, coord)
-        x = self._resize(x)
+        x = tf.saturate_cast(x, tf.uint8)
+        if self.args.dataset == 'imagenet':
+            if self.mode == 'train':
+                x = self._crop(x, shape, coord)
+            else:
+                x = self._centercrop(x, shape)
+
+            x = self._resize(x)
+
+        if self.mode == 'train':
+            x = self._random_color_jitter(x, p=.8)
+            x = self._random_grayscale(x, p=.2)
+            if self.args.dataset == 'imagenet':
+                x = self._random_gaussian_blur(x, p=.5)
+            x = self._random_hflip(x)
+
         x = self._standardize(x)
         return x
 
@@ -46,7 +61,25 @@ def _crop(self, x, shape, coord=[[[0., 0., 1., 1.]]]):
 
         offset_height, offset_width, _ = tf.unstack(bbox_begin)
         target_height, target_width, _ = tf.unstack(bbox_size)
-        x = tf.slice(x, [offset_height, offset_width, 0], [target_height, target_width, 3])
+        x = tf.slice(x, [offset_height, offset_width, 0], [target_height, target_width, -1])
+        return x
+
+    def _centercrop(self, x, shape):
+        if tf.less(shape[0], self.args.img_size):
+            offset_height = 0
+            target_height = shape[0]
+        else:
+            offset_height = tf.maximum(0, shape[0]-self.args.img_size) // 2
+            target_height = self.args.img_size
+
+        if tf.less(shape[1], self.args.img_size):
+            offset_width = 0
+            target_width = shape[1]
+        else:
+            offset_width = tf.maximum(0, shape[1]-self.args.img_size) // 2
+            target_width = self.args.img_size
+
+        x = tf.slice(x, [offset_height, offset_width, 0], [target_height, target_width, -1])
         return x
 
     def _resize(self, x):

callback.py

@@ -1,6 +1,7 @@
 import os
 import six
 import yaml
+import tqdm
 import numpy as np
 import pandas as pd
 import tensorflow as tf
@@ -9,6 +10,8 @@
 from tensorflow.keras.callbacks import CSVLogger
 from tensorflow.keras.callbacks import TensorBoard
 from tensorflow.keras.experimental import CosineDecay
+from sklearn.neighbors import KNeighborsClassifier
+from sklearn.metrics import accuracy_score
 
 from common import create_stamp
 
@@ -90,7 +93,7 @@ def create_callbacks(args, logger, initial_epoch):
                         f'history - {args.history} | '
                         f'tensorboard - {args.tensorboard}')
 
-    callbacks = []        
+    callbacks = []
     if args.checkpoint:
         if args.task == 'pretext':
             callbacks.append(ModelCheckpoint(
@@ -108,13 +111,13 @@ def create_callbacks(args, logger, initial_epoch):
                 save_best_only=True))
         else:
             callbacks.append(ModelCheckpoint(
-                filepath='{args.result_path}/{args.task}/{args.stamp}/checkpoint/latest',
+                filepath=f'{args.result_path}/{args.task}/{args.stamp}/checkpoint/latest',
                 monitor='val_acc1',
                 mode='max',
                 verbose=1,
                 save_weights_only=True))
             callbacks.append(ModelCheckpoint(
-                filepath='{args.result_path}/{args.task}/{args.stamp}/checkpoint/best',
+                filepath=f'{args.result_path}/{args.task}/{args.stamp}/checkpoint/best',
                 monitor='val_acc1',
                 mode='max',
                 verbose=1,

common.py

@@ -21,6 +21,7 @@ def get_arguments():
     parser = argparse.ArgumentParser()
     parser.add_argument("--task",           type=str,       default='pretext',
                         choices=['pretext', 'lincls'])
+    parser.add_argument("--dataset",        type=str,       default='imagenet')
     parser.add_argument("--freeze",         action='store_true')
     parser.add_argument("--backbone",       type=str,       default='resnet50')
     parser.add_argument("--batch_size",     type=int,       default=256)
@@ -38,6 +39,7 @@ def get_arguments():
     parser.add_argument("--steps",          type=int,       default=0)
     parser.add_argument("--epochs",         type=int,       default=200)
 
+    parser.add_argument("--evaluate",       action='store_true')
     parser.add_argument("--checkpoint",     action='store_true')
     parser.add_argument("--history",        action='store_true')
     parser.add_argument("--tensorboard",    action='store_true')
@@ -107,7 +109,7 @@ def create_stamp():
 
 
 def search_same(args):
-    search_ignore = ['checkpoint', 'history', 'tensorboard', 
+    search_ignore = ['evaluate', 'checkpoint', 'history', 'tensorboard', 
                      'tb_interval', 'snapshot', 'summary',
                      'src_path', 'data_path', 'result_path', 
                      'resume', 'stamp', 'gpus', 'ignore_search']
@@ -127,6 +129,9 @@ def search_same(args):
         for k, v in vars(args).items():
             if k in search_ignore:
                 continue
+
+            if k == 'dataset' and k not in desc:
+                desc[k] = 'imagenet'
 
             if v != desc[k]:
                 # if stamp == '210120_Wed_05_19_52':
@@ -157,7 +162,7 @@ def search_same(args):
 
                     if len(ckpt_list) > 0:
                         args.snapshot = f'{args.result_path}/{args.task}/{args.stamp}/checkpoint/{ckpt_list[-1]}'
-                        initial_epoch = int(ckpt_list[-1].split('_')[0])
+                        initial_epoch = int(df['epoch'].iloc[-1]) + 1
                     else:
                         print('{} Training already finished!!!'.format(stamp))
                         return args, -1

dataloader.py

@@ -10,27 +10,32 @@
 AUTO = tf.data.experimental.AUTOTUNE
 
 
-def set_dataset(task, data_path):
-    trainset = pd.read_csv(
-        os.path.join(
-            data_path, 'imagenet_trainset.csv'
-        )).values.tolist()
-    trainset = [[os.path.join(data_path, t[0]), t[1]] for t in trainset]
-
-    if task == 'lincls':
+def set_dataset(task, dataset, data_path):
+    if dataset == 'imagenet':
+        trainset = pd.read_csv(
+            os.path.join(
+                data_path, 'imagenet_trainset.csv'
+            )).values.tolist()
+        trainset = [[os.path.join(data_path, t[0]), t[1]] for t in trainset]
+
         valset = pd.read_csv(
             os.path.join(
                 data_path, 'imagenet_valset.csv'
             )).values.tolist()
         valset = [[os.path.join(data_path, t[0]), t[1]] for t in valset]
-        return np.array(trainset, dtype='object'), np.array(valset, dtype='object')
+        
+    elif dataset == 'cifar10':
+        (x_train, y_train), (x_test, y_test) = tf.keras.datasets.cifar10.load_data()
+        trainset = [[i, l] for i, l in zip(x_train, y_train.flatten())]
+        valset = [[i, l] for i, l in zip(x_test, y_test.flatten())]
 
-    return np.array(trainset, dtype='object')
+    return np.array(trainset, dtype='object'), np.array(valset, dtype='object')
 
 
 class DataLoader:
-    def __init__(self, args, mode, datalist, batch_size, num_workers=1, shuffle=True):
+    def __init__(self, args, task, mode, datalist, batch_size, num_workers=1, shuffle=True):
         self.args = args
+        self.task = task
         self.mode = mode
         self.datalist = datalist
         self.batch_size = batch_size
@@ -50,7 +55,7 @@ def fetch_dataset(self, path, y=None):
         return tf.data.Dataset.from_tensors(x)
 
     def augmentation(self, img, shape):
-        if self.args.task == 'pretext':
+        if self.task == 'pretext':
             img_list = []
             for _ in range(2): # query, key
                 aug_img = tf.identity(img)
@@ -61,20 +66,25 @@ def augmentation(self, img, shape):
             return self.augset._augment_lincls(img, shape)
 
     def dataset_parser(self, value, label=None):
-        shape = tf.image.extract_jpeg_shape(value)
-        img = tf.io.decode_jpeg(value, channels=3)
+        if self.args.dataset == 'imagenet':
+            shape = tf.image.extract_jpeg_shape(value)
+            img = tf.io.decode_jpeg(value, channels=3)
+        elif self.args.dataset == 'cifar10':
+            shape = (32, 32, 3)
+            img = value
+
         if label is None:
-            # moco
+            # pretext
             return self.augmentation(img, shape)
         else:
             # lincls
             inputs = self.augmentation(img, shape)
-            labels = tf.one_hot(label, self.args.classes)
-            return (inputs, labels)
+            # labels = tf.one_hot(label, self.args.classes)
+            return (inputs, label)
 
     def _dataloader(self):
         self.imglist = self.datalist[:,0].tolist()
-        if self.args.task == 'pretext':
+        if self.task == 'pretext':
             dataset = tf.data.Dataset.from_tensor_slices(self.imglist)
         else:
             self.labellist = self.datalist[:,1].tolist()
@@ -84,7 +94,9 @@ def _dataloader(self):
         if self.shuffle:
             dataset = dataset.shuffle(len(self.datalist))
 
-        dataset = dataset.interleave(self.fetch_dataset, num_parallel_calls=AUTO)
+        if self.args.dataset == 'imagenet':
+            dataset = dataset.interleave(self.fetch_dataset, num_parallel_calls=AUTO)
+
         dataset = dataset.map(self.dataset_parser, num_parallel_calls=AUTO)
         dataset = dataset.batch(self.batch_size)
         dataset = dataset.prefetch(AUTO)

layer.py

@@ -1,16 +1,46 @@
 import tensorflow as tf
+from tensorflow.keras.layers import Conv2D
+from tensorflow.keras.layers import Dense
+from tensorflow.keras.layers import BatchNormalization
+from tensorflow.keras.layers.experimental import SyncBatchNormalization
 from tensorflow.keras.initializers import Constant
 
 
-class Conv2D(tf.keras.layers.Conv2D):
-    def build(self, input_shape):
-        k = 1 / input_shape[-1]
-        self.kernel_initializer = Constant(tf.random.uniform([], -tf.sqrt(k), tf.sqrt(k)))
-        super(Conv2D, self).build(input_shape)
+BatchNorm_DICT = {
+    "bn": BatchNormalization,
+    "syncbn": SyncBatchNormalization}
 
 
-class Dense(tf.keras.layers.Dense):
-    def build(self, input_shape):
-        k = 1 / input_shape[-1]
-        self.kernel_initializer = Constant(tf.random.uniform([], -tf.sqrt(k), tf.sqrt(k)))
-        super(Dense, self).build(input_shape)
+def _conv2d(**custom_kwargs):
+    def _func(*args, **kwargs):
+        kwargs.update(**custom_kwargs)
+        return Conv2D(*args, **kwargs)
+    return _func
+
+
+def _batchnorm(norm='bn', **custom_kwargs):
+    def _func(*args, **kwargs):
+        kwargs.update(**custom_kwargs)
+        return BatchNorm_DICT[norm](*args, **kwargs)
+    return _func
+
+
+def _dense(**custom_kwargs):
+    def _func(*args, **kwargs):
+        kwargs.update(**custom_kwargs)
+        return Dense(*args, **kwargs)
+    return _func
+
+
+# class Conv2D(tf.keras.layers.Conv2D):
+#     def build(self, input_shape):
+#         k = 1 / input_shape[-1]
+#         self.kernel_initializer = Constant(tf.random.uniform([], -tf.sqrt(k), tf.sqrt(k)))
+#         super(Conv2D, self).build(input_shape)
+
+
+# class Dense(tf.keras.layers.Dense):
+#     def build(self, input_shape):
+#         k = 1 / input_shape[-1]
+#         self.kernel_initializer = Constant(tf.random.uniform([], -tf.sqrt(k), tf.sqrt(k)))
+#         super(Dense, self).build(input_shape)

main.py

@@ -21,14 +21,17 @@ def train_pretext(args, logger, initial_epoch, strategy, num_workers):
     ##########################
     # Dataset
     ##########################
-    trainset = set_dataset(args.task, args.data_path)
+    trainset, valset = set_dataset(args.task, args.dataset, args.data_path)
     steps_per_epoch = args.steps or len(trainset) // args.batch_size
 
     logger.info("TOTAL STEPS OF DATASET FOR TRAINING")
     logger.info("========== TRAINSET ==========")
     logger.info(f"    --> {len(trainset)}")
     logger.info(f"    --> {steps_per_epoch}")
 
+    logger.info("=========== VALSET ===========")
+    logger.info(f"    --> {len(valset)}")
+
 
     ##########################
     # Model & Generator
@@ -41,8 +44,9 @@ def train_pretext(args, logger, initial_epoch, strategy, num_workers):
             optimizer=tf.keras.optimizers.SGD(lr_scheduler, momentum=.9),
             loss=tf.keras.losses.cosine_similarity,
             run_eagerly=False)
-    
-    train_generator = DataLoader(args, 'train', trainset, args.batch_size, num_workers).dataloader
+
+    train_generator = DataLoader(args, args.task, 'train', trainset, args.batch_size, num_workers).dataloader
+
 
     ##########################
     # Train
@@ -63,11 +67,11 @@ def train_pretext(args, logger, initial_epoch, strategy, num_workers):
 
 
 def train_lincls(args, logger, initial_epoch, strategy, num_workers):
-    assert args.snapshot is not None, 'pretrained weight is needed!'
+    # assert args.snapshot is not None, 'pretrained weight is needed!'
     ##########################
     # Dataset
     ##########################
-    trainset, valset = set_dataset(args.task, args.data_path)
+    trainset, valset = set_dataset(args.task, args.dataset, args.data_path)
     steps_per_epoch = args.steps or len(trainset) // args.batch_size
     validation_steps = len(valset) // args.batch_size
 
@@ -84,19 +88,20 @@ def train_lincls(args, logger, initial_epoch, strategy, num_workers):
     ##########################
     # Model & Generator
     ##########################
+    train_generator = DataLoader(args, args.task, 'train', trainset, args.batch_size, num_workers).dataloader
+    val_generator = DataLoader(args, args.task, 'val', valset, args.batch_size, num_workers).dataloader
+        
     with strategy.scope():
         backbone = SimSiam(args, logger)
         model = set_lincls(args, backbone.encoder)
 
         lr_scheduler = OptionalLearningRateSchedule(args, steps_per_epoch, initial_epoch)
         model.compile(
             optimizer=tf.keras.optimizers.SGD(lr_scheduler, momentum=.9),
-            metrics=[tf.keras.metrics.TopKCategoricalAccuracy(1, 'acc1', dtype=tf.float32),
-                     tf.keras.metrics.TopKCategoricalAccuracy(5, 'acc5', dtype=tf.float32)],
-            loss=tf.keras.losses.CategoricalCrossentropy(from_logits=True, name='loss'))
-
-    train_generator = DataLoader(args, 'train', trainset, args.batch_size, num_workers).dataloader
-    val_generator = DataLoader(args, 'val', valset, args.batch_size, num_workers).dataloader
+            metrics=[tf.keras.metrics.SparseTopKCategoricalAccuracy(1, 'acc1', dtype=tf.float32),
+                     tf.keras.metrics.SparseTopKCategoricalAccuracy(5, 'acc5', dtype=tf.float32)],
+            loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True, name='loss'),
+            run_eagerly=False)
 
 
     ##########################
@@ -123,9 +128,15 @@ def main():
     set_seed()
     args = get_arguments()
     if args.task == 'pretext':
-        args.lr = 0.5 * float(args.batch_size / 256)
+        if args.dataset == 'imagenet':
+            args.lr = 0.5 * float(args.batch_size / 256)
+        elif args.dataset == 'cifar10':
+            args.lr = 0.03 * float(args.batch_size / 256)
     else:
-        args.lr = 30. * float(args.batch_size / 256)
+        if args.dataset == 'imagenet' and args.freeze:
+            args.lr = 30. * float(args.batch_size / 256)
+        else:# args.dataset == 'cifar10':
+            args.lr = 1.8 * float(args.batch_size / 256)
 
     args, initial_epoch = search_same(args)
     if initial_epoch == -1: