예제 #1
0
def main(model_num=1):

    preprocess = Preprocess()

    texts_train, labels_train = preprocess.preprocessData(
        '../projet2/train.txt', mode="train")
    texts_dev, labels_dev = preprocess.preprocessData('../projet2/dev.txt',
                                                      mode="train")

    MAX_SEQUENCE_LENGTH = 24
    LSTM_DIM = 64
    HIDDEN_LAYER_DIM = 30
    NUM_CLASSES = 4
    GAUSSIAN_NOISE = 0.1
    DROPOUT = 0.2
    DROPOUT_LSTM = 0.2
    BATCH_SIZE = 200

    X_train, X_val, y_train, y_val = train_test_split(texts_train,
                                                      labels_train,
                                                      test_size=0.2,
                                                      random_state=42)

    labels_categorical_train = to_categorical(np.asarray(y_train))
    labels_categorical_val = to_categorical(np.asarray(y_val))
    labels_categorical_dev = to_categorical(np.asarray(labels_dev))

    embedding = Embedding('../projet2/emosense.300d.txt')
    embeddings = embedding.getMatrix()
    tokenizer = embedding.getTokenizer()

    message_first_message_train, message_second_message_train, message_third_message_train = get_sequences(
        X_train, MAX_SEQUENCE_LENGTH, tokenizer)
    message_first_message_val, message_second_message_val, message_third_message_val = get_sequences(
        X_val, MAX_SEQUENCE_LENGTH, tokenizer)
    message_first_message_dev, message_second_message_dev, message_third_message_dev = get_sequences(
        texts_dev, MAX_SEQUENCE_LENGTH, tokenizer)

    model = CustomModel(model_num)
    model.build(embeddings,
                MAX_SEQUENCE_LENGTH,
                LSTM_DIM,
                HIDDEN_LAYER_DIM,
                NUM_CLASSES,
                noise=GAUSSIAN_NOISE,
                dropout_lstm=DROPOUT_LSTM,
                dropout=DROPOUT)
    model.summary()
    history = model.train(message_first_message_train,
                          message_second_message_train,
                          message_third_message_train,
                          labels_categorical_train, message_first_message_val,
                          message_second_message_val,
                          message_third_message_val, labels_categorical_val)

    y_pred = model.predict([
        message_first_message_dev, message_second_message_dev,
        message_third_message_dev
    ])
예제 #2
0
    def _init_learner(self):
        self.cNNModel = CustomModel()

        self._init_nodes()
        self._init_predictions()

        # We'll compute them only once in a while by calling their {eval()} method.
        #self.train_all_prediction = tf.nn.softmax(self.cNNModel.model_func()(self.train_all_data_node))

        self._init_regularizer()
        self._init_loss()
        self._init_learning_rate()
        self._init_optimizer()
        self._init_metrics()

        #self._init_params_summaries()

        # Add ops to save and restore all the variables.
        self.saver = tf.train.Saver()
예제 #3
0
def save_hard_images(
        threshold_list=(0.9999, 0.999, 0.99, 0.97, 0.95, 0.9, 0.8, 0.7, 0.5),
        image_number=1250):

    CustomModel = import_model()
    model = CustomModel()
    factor = 1
    for i, threshold in enumerate(threshold_list):
        print("acquisition::get_hard_images : threshold :", threshold)
        try:

            shutil.rmtree(
                os.path.join(negatives_hard_directory, str(threshold)))
        except FileNotFoundError:
            pass
        os.makedirs(os.path.join(negatives_hard_directory, str(threshold)))

        model.train(image_number, image_number)
        pyramid = Pyramid(model, threshold)
        pyramid.add_false_positive_to_negative_db(directory_path=added_images,
                                                  save_path=os.path.join(
                                                      negatives_hard_directory,
                                                      str(threshold)),
                                                  strides=(9, 9))

        pyramid = Pyramid(model, threshold_list[max(0, i - 1)])
        pyramid.test_pyramide('photo_famille.jpeg')

        factor *= 0.8
        model.recompile(factor)
예제 #4
0
def init_inference(model_kwargs):
    global model
    global device
    model = CustomModel(**model_kwargs)
    model.eval()

    if args.trt_module :
        from torch2trt import TRTModule
        if args.trt_conversion :
            model.load_state_dict(torch.load(args.pretrained_model))
            model = model.cuda()
            x = torch.ones((1, 3, 240, 320)).cuda()
            from torch2trt import torch2trt
            model_trt = torch2trt(model, [x], max_batch_size=100, fp16_mode=True)
            #model_trt = torch2trt(model, [x], max_batch_size=100)
            torch.save(model_trt.state_dict(), args.trt_model)
            exit()
        model_trt = TRTModule()
        #model_trt.load_state_dict(torch.load('road_following_model_trt_half.pth'))
        model_trt.load_state_dict(torch.load(args.trt_model))
        model = model_trt.to(device)
    else :
        model.load_state_dict(torch.load(args.pretrained_model))
        model = model.to(device)
예제 #5
0
    print(f'\nwhat else does the wrapper see...\n')
    print(mgr.proj_cfg)
    print(mgr.basic_model_cfg)
    print(mgr.tune_model_cfg)
    print(mgr.build_id)
    print(mgr.artifact_path)
    print(mgr.args.predict_one)
    print(f'\n')

    time.sleep(1)
    #----

    # CALL MODEL
    como = CustomModel(mgr.proj_cfg, mgr.basic_model_cfg, mgr.tune_model_cfg,
                       mgr.build_id, mgr.artifact_path, mgr.train_dir,
                       mgr.test_dir)

    print(f'\nwhat does the model see...\n')
    como.debug_dump()

    #como.load_basic_model()
    #como.setup_basic_model()
    #como.run_model()

    #como.save_model()

    # -- OR --
    # como.load_custom_model(filepath)
    # como.setup_custom_model()
    # como.run_model()
            net_out = Dropout(hypes["dropout"])(net_out)
        classifier = Dense(units=hypes["num_classes"],
                           activation=hypes["classifier"])(net_out)
        final_model = Model(inputs=pretrained.input, outputs=classifier)
    else:
        final_model = Model(inputs=pretrained.input, outputs=classifier)

    final_model.compile(loss=keras.losses.sparse_categorical_crossentropy,
                        optimizer=keras.optimizers.Adam(
                            lr=hypes["learning_rate"],
                            decay=hypes["weight_decay"]),
                        metrics=hypes["metric"])
    print(final_model.summary())

    #fit network
    CustomModel.train_model(final_model, hypes, training_generator,
                            validation_generator)

    #inference
    hypes_test = hypes
    hypes_test["batch_size"] = 1
    testing_generator = DataGenerator(testing_set[index],
                                      classes,
                                      hypes_test,
                                      training=False)
    model_performance = {}
    for filename in os.listdir(hypes["chkp_dir"] + hypes["version"]):
        if "weights" in filename:
            model_performance[int(float(filename[-7:-3]) * 100)] = filename
    final_model.load_weights(hypes["chkp_dir"] + hypes["version"] +
                             model_performance[max(model_performance.keys())])
    _, _, roc, _ = CustomModel.test_model(final_model, hypes_test,
예제 #7
0
def main():
    device = "cuda:0" if torch.cuda.is_available() else "cpu"

    parser = argparse.ArgumentParser()
    parser.add_argument('--image_path', type=str, default="./data/cache/train")
    parser.add_argument('--label_path', type=str, default="./data/cache/train.csv")
    parser.add_argument('--kfold_idx', type=int, default=0)

    # parser.add_argument('--model', type=str, default='CustomModel')
    parser.add_argument('--model', type=str, default='efficientnet-b0')
    parser.add_argument('--epochs', type=int, default=2000)
    parser.add_argument('--batch_size', type=int, default=50)
    parser.add_argument('--lr', type=float, default=1e-3)
    parser.add_argument('--patient', type=int, default=8)
    parser.add_argument('--seed', type=int, default=42)

    parser.add_argument('--device', type=str, default=device)
    parser.add_argument('--resume', type=str, default=None)
    parser.add_argument('--comments', type=str, default=None)

    args = parser.parse_args()

    print('=' * 50)
    print('[info msg] arguments\n')
    for key, value in vars(args).items():
        print(key, ":", value)
    print('=' * 50)
    
    assert os.path.isdir(args.image_path), 'wrong path'
    assert os.path.isfile(args.label_path), 'wrong path'
    if (args.resume):
        assert os.path.isfile(args.resume), 'wrong path'
    # assert args.kfold_idx < 5

    seed_everything(args.seed)

    data_df = pd.read_csv(args.label_path)

    sss = StratifiedShuffleSplit(n_splits=1, test_size=0.3, random_state=args.seed)
    for train_idx, valid_idx in sss.split(X=data_df['id'], y=data_df['accent']):
        train_df = data_df.iloc[train_idx]
        valid_df = data_df.iloc[valid_idx]

    train_data = dataset.DaconDataset(
        image_folder=args.image_path,
        label_df=train_df,
     )
    
    valid_data = dataset.DaconDataset(
        image_folder=args.image_path,
        label_df=valid_df,
    )

    train_sampler = get_sampler(
        df=train_df,
        dataset=train_data
    )

    valid_sampler = get_sampler(
        df=valid_df,
        dataset=valid_data
    )

    train_data_loader = torch.utils.data.DataLoader(
            train_data,
            batch_size=args.batch_size,
            # shuffle=True,
            sampler=train_sampler
        )

    valid_data_loader = torch.utils.data.DataLoader(
            valid_data,
            batch_size=args.batch_size,
            # shuffle=False,
            sampler=valid_sampler
        )

    model = None

    if args.model == 'CustomModel':
        model = CustomModel()
        print('[info msg] {} model is created\n'.format('CustomModel'))
    else:
        model = EfficientNet.from_pretrained(args.model, in_channels=1, num_classes=6, dropout_rate=0.3, advprop=True)
        print('[info msg] {} model is created\n'.format(args.model))
    
    print('=' * 50)

    if(args.resume):
        model.load_state_dict(torch.load(args.resume))
        print('[info msg] pre-trained weight is loaded !!\n')        
        print(args.resume)
        print('=' * 50)

    if args.device == 'cuda' and torch.cuda.device_count() > 1 :
        model = torch.nn.DataParallel(model)
 
    ##### Wandb ######
    wandb.init(project='dacon_voice')
    wandb.run.name = args.comments
    wandb.config.update(args)
    wandb.watch(model)
    ##################
    
    model.to(args.device)

    optimizer = torch.optim.Adam(model.parameters(), args.lr)
    criterion = torch.nn.CrossEntropyLoss()
    scheduler = ReduceLROnPlateau(
        optimizer=optimizer,
        mode='min',
        patience=2,
        factor=0.5,
        verbose=True
        )

    train_loss = []
    train_acc = []
    valid_loss = []
    valid_acc = []

    best_loss = float("inf")

    patient = 0

    date_time = datetime.now().strftime("%m%d%H%M%S")
    SAVE_DIR = os.path.join('./model', date_time)

    print('[info msg] training start !!\n')
    startTime = datetime.now()
    for epoch in range(args.epochs):        
        print('Epoch {}/{}'.format(epoch+1, args.epochs))
        train_epoch_loss, train_epoch_acc = trainer.train(
            train_loader=train_data_loader,
            model=model,
            loss_func=criterion,
            device=args.device,
            optimizer=optimizer,
            )
        train_loss.append(train_epoch_loss)
        train_acc.append(train_epoch_acc)

        valid_epoch_loss, valid_epoch_acc = trainer.validate(
            valid_loader=valid_data_loader,
            model=model,
            loss_func=criterion,
            device=args.device,
            scheduler=scheduler,
            )
        valid_loss.append(valid_epoch_loss)        
        valid_acc.append(valid_epoch_acc)

        wandb.log({
            "Train Acc": train_epoch_acc,
            "Valid Acc": valid_epoch_acc,
            "Train Loss": train_epoch_loss,
            "Valid Loss": valid_epoch_loss,
            })

        if best_loss > valid_epoch_loss:
            patient = 0
            best_loss = valid_epoch_loss

            Path(SAVE_DIR).mkdir(parents=True, exist_ok=True)
            torch.save(model.state_dict(), os.path.join(SAVE_DIR, 'model_best.pth'))
            print('MODEL IS SAVED TO {}!!!'.format(date_time))
            
        else:
            patient += 1
            if patient > args.patient - 1:
                print('=======' * 10)
                print("[Info message] Early stopper is activated")
                break

    elapsed_time = datetime.now() - startTime

    train_loss = np.array(train_loss)
    train_acc = np.array(train_acc)
    valid_loss = np.array(valid_loss)
    valid_acc = np.array(valid_acc)

    best_loss_pos = np.argmin(valid_loss)
    
    print('=' * 50)
    print('[info msg] training is done\n')
    print("Time taken: {}".format(elapsed_time))
    print("best loss is {} w/ acc {} at epoch : {}".format(best_loss, valid_acc[best_loss_pos], best_loss_pos))    

    print('=' * 50)
    print('[info msg] {} model weight and log is save to {}\n'.format(args.model, SAVE_DIR))

    with open(os.path.join(SAVE_DIR, 'log.txt'), 'w') as f:
        for key, value in vars(args).items():
            f.write('{} : {}\n'.format(key, value))            

        f.write('\n')
        f.write('total ecpochs : {}\n'.format(str(train_loss.shape[0])))
        f.write('time taken : {}\n'.format(str(elapsed_time)))
        f.write('best_train_loss {} w/ acc {} at epoch : {}\n'.format(np.min(train_loss), train_acc[np.argmin(train_loss)], np.argmin(train_loss)))
        f.write('best_valid_loss {} w/ acc {} at epoch : {}\n'.format(np.min(valid_loss), valid_acc[np.argmin(valid_loss)], np.argmin(valid_loss)))

    plt.figure(figsize=(15,5))
    plt.subplot(1, 2, 1)
    plt.plot(train_loss, label='train loss')
    plt.plot(valid_loss, 'o', label='valid loss')
    plt.axvline(x=best_loss_pos, color='r', linestyle='--', linewidth=1.5)
    plt.legend()
    plt.subplot(1, 2, 2)
    plt.plot(train_acc, label='train acc')
    plt.plot(valid_acc, 'o', label='valid acc')
    plt.axvline(x=best_loss_pos, color='r', linestyle='--', linewidth=1.5)
    plt.legend()
    plt.savefig(os.path.join(SAVE_DIR, 'history.png'))
예제 #8
0
                # statistics
                running_loss += loss.data[0] * inputs.size(0)
                running_corrects += torch.sum(preds == labels.data)

            epoch_loss = running_loss / len(dataloaders[phase])
            epoch_acc = running_corrects / len(dataloaders[phase])

            print('{} Loss: {:.4f} Acc: {:.4f}'.format(
                phase, epoch_loss, epoch_acc))

            # deep copy the model
            if phase == 'val' and epoch_acc > best_acc:
                best_acc = epoch_acc
                best_model_wts = copy.deepcopy(model.state_dict())
                torch.save(best_model_wts, "models2/model-%s.weights" % epoch)

    print('Training complete')
    print('Best val Acc: {:4f}'.format(best_acc))

if __name__=='__main__': 
    num_classes = 3
    model = CustomModel()

    train_path = "dataset2"
    test_path = "test_set"
    train_loader = get_dataloader(train_path, batch_size=8)
    test_loader = get_dataloader(test_path, batch_size=30)

    loss = nn.CrossEntropyLoss()
    train(model, loss, train_loader, test_loader)
예제 #9
0
from model import CustomModel
from data import transform_driving_image, LEFT, RIGHT, GO

id_to_steer = {
    LEFT: -1,
    RIGHT: 1,
    GO: 0,
}

if __name__ == '__main__':

    if len(sys.argv) < 2:
        sys.exit("Usage : python drive.py path/to/weights")
    # load the model
    model_weights = sys.argv[1]
    model = CustomModel()
    model.load_state_dict(torch.load(model_weights))

    env = gym.make('CarRacing-v0').env
    env.reset()

    a = np.array([0.0, 0.0, 0.0])

    def key_press(k, mod):
        global restart
        if k == key.LEFT: a[0] = -1.0
        if k == key.RIGHT: a[0] = +1.0
        if k == key.UP: a[1] = +1.0
        if k == key.DOWN:
            a[2] = +0.8  # set 1.0 for wheels to block to zero rotation
예제 #10
0
	tr = transforms.Compose([transforms.RandomResizedCrop(300),
                             transforms.ToTensor(),
                             transforms.Normalize([0.4589, 0.4355, 0.4032],[0.2239, 0.2186, 0.2206])])

    augs = transforms.Compose([transforms.RandomResizedCrop(300),
                               transforms.RandomRotation(20),
                               transforms.ToTensor(),
                               transforms.Normalize([0.4589, 0.4355, 0.4032],[0.2239, 0.2186, 0.2206])])

    train_set = CustomDataset(dir_csv, dir_img, transforms=augs)
    train_loader = DataLoader(train_set, batch_size=batch_size_train, shuffle=True)

    val_set = CustomDataset(dir_csv, dir_img, transforms=tr)
    val_loader = DataLoader(val_set, batch_size=batch_size_test, shuffle=False)

    model = CustomModel()
    loss_function = CustomLoss()

    model.to(device)
    print('Starting optimizer with LR={}'.format(lr))
    optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9)
   
    scheduler = StepLR(optimizer, step_size=1, gamma=gamma)
    for epoch in range(1, num_epochs + 1):
        train(model, device, train_loader, optimizer, epoch, loss_function)
        test(model, device, test_loader, loss_function)
        scheduler.step()

    torch.save(model.state_dict(), "well_trained model.pt")

예제 #11
0
class Learner:
    def __init__(self):
        self._init_learner()

    def _init_learner(self):
        self.cNNModel = CustomModel()

        self._init_nodes()
        self._init_predictions()

        # We'll compute them only once in a while by calling their {eval()} method.
        #self.train_all_prediction = tf.nn.softmax(self.cNNModel.model_func()(self.train_all_data_node))

        self._init_regularizer()
        self._init_loss()
        self._init_learning_rate()
        self._init_optimizer()
        self._init_metrics()

        #self._init_params_summaries()

        # Add ops to save and restore all the variables.
        self.saver = tf.train.Saver()

    def _init_nodes(self):
        self.data_node = tf.placeholder(tf.float32,
                                        shape=(None, EFFECTIVE_INPUT_SIZE,
                                               EFFECTIVE_INPUT_SIZE,
                                               NUM_CHANNELS))
        self.labels_node = tf.placeholder(tf.float32, shape=(None, NUM_LABELS))

    def _init_regularizer(self):
        # L2 regularization for the fully connected parameters.
        self.regularizers = self.cNNModel.get_weights()

    def _init_learning_rate(self):
        # Optimizer: set up a variable that's incremented once per batch and
        # controls the learning rate decay.
        self.batch = tf.Variable(0)

    def _init_optimizer(self):
        # Use simple momentum for the optimization.
        self.optimizer = tf.train.AdamOptimizer(
            ADAM_INITIAL_LEARNING_RATE).minimize(self.loss,
                                                 global_step=self.batch)

    def _init_predictions(self):
        self.logits = self.cNNModel.model_func()(self.data_node, True)
        self.logits_validation = self.cNNModel.model_func()(self.data_node,
                                                            False)
        self.logits_test = self.cNNModel.model_func()(self.data_node, False)

        # Predictions for the minibatch, validation set and test set.
        self.train_predictions = tf.nn.softmax(self.logits)
        self.validation_predictions = tf.nn.softmax(self.logits_validation)
        self.test_predictions = tf.nn.softmax(self.logits_test)

    def _init_loss(self):
        # print 'logits = ' + str(logits.get_shape()) + ' train_labels_node = ' + str(train_labels_node.get_shape())
        self.loss = tf.reduce_mean(
            tf.nn.softmax_cross_entropy_with_logits(logits=self.logits,
                                                    labels=self.labels_node))

        # Add the regularization term to the loss.
        if USE_L2_REGULARIZATION:
            self.loss += 5e-4 * self.regularizers

    def _init_metrics(self):
        l_train = tf.argmax(self.labels_node, 1)
        p_train = tf.argmax(self.train_predictions, 1)

        l_validation = tf.argmax(self.labels_node, 1)
        p_validation = tf.argmax(self.validation_predictions, 1)

        l_test = tf.argmax(self.labels_node, 1)
        p_test = tf.argmax(self.validation_predictions, 1)
        """
        TRAINING METRICS
        """
        self.true_train_pos, self.true_train_pos_op = tf.contrib.metrics.streaming_true_positives(
            labels=l_train, predictions=p_train, name="train")

        self.false_train_pos, self.false_train_pos_op = tf.contrib.metrics.streaming_false_positives(
            labels=l_train, predictions=p_train, name="train")

        self.true_train_neg, self.true_train_neg_op = tf.contrib.metrics.streaming_true_negatives(
            labels=l_train, predictions=p_train, name="train")

        self.false_train_neg, self.false_train_neg_op = tf.contrib.metrics.streaming_false_negatives(
            labels=l_train, predictions=p_train, name="train")
        """
        VALIDATION METRICS
        """
        self.true_validation_pos, self.true_validation_pos_op = tf.contrib.metrics.streaming_true_positives(
            labels=l_validation, predictions=p_validation, name="validation")

        self.false_validation_pos, self.false_validation_pos_op = tf.contrib.metrics.streaming_false_positives(
            labels=l_validation, predictions=p_validation, name="validation")

        self.true_validation_neg, self.true_validation_neg_op = tf.contrib.metrics.streaming_true_negatives(
            labels=l_validation, predictions=p_validation, name="validation")

        self.false_validation_neg, self.false_validation_neg_op = tf.contrib.metrics.streaming_false_negatives(
            labels=l_validation, predictions=p_validation, name="validation")
        """
        TEST METRICS
        """
        self.true_test_pos, self.true_test_pos_op = tf.contrib.metrics.streaming_true_positives(
            labels=l_test, predictions=p_test, name="test")

        self.false_test_pos, self.false_test_pos_op = tf.contrib.metrics.streaming_false_positives(
            labels=l_test, predictions=p_test, name="test")

        self.true_test_neg, self.true_test_neg_op = tf.contrib.metrics.streaming_true_negatives(
            labels=l_test, predictions=p_test, name="test")

        self.false_test_neg, self.false_test_neg_op = tf.contrib.metrics.streaming_false_negatives(
            labels=l_test, predictions=p_test, name="test")

    def model(self):
        return self.cNNModel

    def update_feed_dictionary(self, batch_data, batch_labels):
        self.feed_dictionary = {
            self.data_node: batch_data,
            self.labels_node: batch_labels
        }

    def get_feed_dictionnary(self):
        return self.feed_dictionary

    def get_train_ops(self):
        return [self.optimizer, self.loss, self.train_predictions]

    def get_train_metric_ops(self):
        return [
            self.true_train_pos, self.false_train_pos, self.true_train_neg,
            self.false_train_neg
        ]

    def get_train_metric_update_ops(self):
        return [
            self.true_train_pos_op, self.false_train_pos_op,
            self.true_train_neg_op, self.false_train_neg_op
        ]

    def get_validation_ops(self):
        return [self.validation_predictions]

    def get_validation_metric_ops(self):
        return [
            self.true_validation_pos, self.false_validation_pos,
            self.true_validation_neg, self.false_validation_neg
        ]

    def get_validation_metric_update_ops(self):
        return [
            self.true_validation_pos_op, self.false_validation_pos_op,
            self.true_validation_neg_op, self.false_validation_neg_op
        ]

    def get_test_ops(self):
        return [self.test_predictions]

    def get_test_metric_ops(self):
        return [
            self.true_test_pos, self.false_test_pos, self.true_test_neg,
            self.false_test_neg
        ]

    def get_test_metric_update_ops(self):
        return [
            self.true_test_pos_op, self.false_test_pos_op,
            self.true_test_neg_op, self.false_test_neg_op
        ]
예제 #12
0
image_res = (240, 320)

data_transforms = transforms.Compose([
    transforms.Resize(image_res),
    transforms.ToTensor()
])  # transforms to resize and convert the image to a tensor object

data_dir = 'images'  # folders with all the images separated by class in subfolders
dataset = datasets.ImageFolder(data_dir, data_transforms)
img_loader = DataLoader(
    dataset, batch_size=4,
    shuffle=True)  # separates the dataset in batches for training
# the batch size may be limited by memory. Reduce it if that is the case

n_classes = 2  # classify 2 different images
net = CustomModel(image_res,
                  n_classes).to(device)  # custom model initialization

optimizer = optim.Adam(net.parameters(),
                       lr=0.001)  # optimizer (basically a gradient descent)

loss = torch.nn.CrossEntropyLoss()  # loss function for classification

for epoch in range(n_epochs):

    epoch_loss = 0.0
    print(epoch)
    for i, data in enumerate(img_loader,
                             0):  # iterator to get the batches for training

        inputs, labels = data
예제 #13
0
            # subset="validation",
            color_mode="grayscale")
        '''
        test_generator = test_datagen.flow_from_directory(
            testdata_dir,
            target_size=(HEIGHT, WIDTH),
            batch_size=1,
            class_mode="categorical",
            subset="validation",
            color_mode="grayscale")
        '''

        model_name = 'test1'

        model_init = CustomModel(input_shape=(HEIGHT, WIDTH, 1),
                                 num_classes=num_classes,
                                 alpha=1.0)
        model = model_init.make_model()

        print(model.summary())

        # SGD_opt = keras.optimizers.SGD(lr=0.01, momentum=0.0, decay=0.1, nesterov=False)
        # Admas_opt = keras.optimizers.Adam(lr=0.001, beta_1=0.9, beta_2=0.999, epsilon=None, decay=0.0, amsgrad=False)
        model.compile(loss=keras.losses.categorical_crossentropy,
                      optimizer='SGD',
                      metrics=['accuracy'])

        filepath = "save/%s.h5" % model_name

        checkpoint = ModelCheckpoint(filepath,
                                     monitor='val_acc',