def main(args): dataset = args.dataset data_type = args.data_type # data dir, image size and batch size DATA_DIR = 'data' TRAIN_DIR = os.path.join(DATA_DIR, 'train') VALID_DIR = os.path.join(DATA_DIR, 'valid') SIZE = (224, 224) BATCH_SIZE = 4 # remove files try: shutil.rmtree(DATA_DIR) except: pass train_ratio = 0.8 # 80% data for training, the rest for testing # get the list of filenames and corresponding list of labels for training et validation train_filenames = [] train_labels = [] val_filenames = [] val_labels = [] # read files into label and frame lists with open('../data/labels/' + dataset + '_' + data_type + '_label.csv') as f: frames_labels = [(line.strip().split(',')[0], line.strip().split(',')[1]) for line in f] # re-organize data by labels file_dir = '../data/' + dataset + '/' + data_type + '/jpg/' file_format = '.jpeg' dict_frame = { } # key: label value, value: a list of indice in the original file for fr_lb in frames_labels: fr, lb = fr_lb if (lb not in dict_frame): dict_frame[lb] = [] dict_frame[lb].append(file_dir + fr + file_format) random.seed() # using current time as the seed # generate filenames and labels for training and validation dataset for lb in dict_frame: # pick random indices for training data for lb in dict_frame train_index = random.sample(range(0, len(dict_frame[lb])), int(train_ratio * len(dict_frame[lb]))) for index in range(len(dict_frame[lb])): # training data if (index in train_index): train_filenames.append(dict_frame[lb][index]) train_labels.append(int(lb) - 1) # validation data else: val_filenames.append(dict_frame[lb][index]) val_labels.append(int(lb) - 1) assert set(train_labels) == set( val_labels), "Train and val labels don't correspond:\n{}\n{}".format( set(train_labels), set(val_labels)) # create new dir data/train/label_x and data/valid/label_x for label in set(train_labels): os.makedirs(os.path.join(TRAIN_DIR, str(label)), exist_ok=True) os.makedirs(os.path.join(VALID_DIR, str(label)), exist_ok=True) # copy files for tr_file, label in zip(train_filenames, train_labels): shutil.copy2(tr_file, os.path.join(TRAIN_DIR, str(label))) for val_file, label in zip(val_filenames, val_labels): shutil.copy2(val_file, os.path.join(VALID_DIR, str(label))) # train models num_train_samples = sum([len(files) for r, d, files in os.walk(TRAIN_DIR)]) num_valid_samples = sum([len(files) for r, d, files in os.walk(VALID_DIR)]) num_train_steps = math.floor(num_train_samples / BATCH_SIZE) num_valid_steps = math.floor(num_valid_samples / BATCH_SIZE) gen = image.ImageDataGenerator() val_gen = image.ImageDataGenerator(horizontal_flip=True, vertical_flip=True) batches = gen.flow_from_directory(TRAIN_DIR, target_size=SIZE, class_mode='categorical', shuffle=True, batch_size=BATCH_SIZE) val_batches = val_gen.flow_from_directory(VALID_DIR, target_size=SIZE, class_mode='categorical', shuffle=True, batch_size=BATCH_SIZE) model = DenseNet121() classes = list(iter(batches.class_indices)) model.layers.pop() for layer in model.layers: layer.trainable = False last = model.layers[-1].output x = Dense(len(classes), activation="softmax")(last) finetuned_model = Model(model.input, x) finetuned_model.compile(optimizer=Adam(lr=0.0001), loss='categorical_crossentropy', metrics=['accuracy']) for c in batches.class_indices: classes[batches.class_indices[c]] = c finetuned_model.classes = classes early_stopping = EarlyStopping(patience=450) checkpointer = ModelCheckpoint('./densenet_model/densenet_121_best.h5', verbose=1, save_best_only=True) finetuned_model.fit_generator(batches, steps_per_epoch=num_train_steps, epochs=450, callbacks=[early_stopping, checkpointer], validation_data=val_batches, validation_steps=num_valid_steps) finetuned_model.save('./densenet_model/densenet_121_final.h5')
def efficientnet_model(channels, init_block_channels, final_block_channels, kernel_sizes, strides_per_stage, expansion_factors, dropout_rate=0.2, tf_mode=False, bn_epsilon=1e-5, in_channels=3, in_size=(224, 224), classes=1000): """ EfficientNet(-B0) model from 'EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks,' https://arxiv.org/abs/1905.11946. Parameters: ---------- channels : list of list of int Number of output channels for each unit. init_block_channels : list of 2 int Numbers of output channels for the initial unit. final_block_channels : int Number of output channels for the final block of the feature extractor. kernel_sizes : list of list of int Number of kernel sizes for each unit. strides_per_stage : list int Stride value for the first unit of each stage. expansion_factors : list of list of int Number of expansion factors for each unit. dropout_rate : float, default 0.2 Fraction of the input units to drop. Must be a number between 0 and 1. tf_mode : bool, default False Whether to use TF-like mode. bn_epsilon : float, default 1e-5 Small float added to variance in Batch norm. in_channels : int, default 3 Number of input channels. in_size : tuple of two ints, default (224, 224) Spatial size of the expected input image. classes : int, default 1000 Number of classification classes. """ input_shape = (in_channels, in_size[0], in_size[1]) if is_channels_first() else\ (in_size[0], in_size[1], in_channels) input = nn.Input(shape=input_shape) activation = "swish" x = effi_init_block(x=input, in_channels=in_channels, out_channels=init_block_channels, bn_epsilon=bn_epsilon, activation=activation, tf_mode=tf_mode, name="features/init_block") in_channels = init_block_channels for i, channels_per_stage in enumerate(channels): kernel_sizes_per_stage = kernel_sizes[i] expansion_factors_per_stage = expansion_factors[i] for j, out_channels in enumerate(channels_per_stage): kernel_size = kernel_sizes_per_stage[j] expansion_factor = expansion_factors_per_stage[j] strides = strides_per_stage[i] if (j == 0) else 1 if i == 0: x = effi_dws_conv_unit(x=x, in_channels=in_channels, out_channels=out_channels, strides=strides, bn_epsilon=bn_epsilon, activation=activation, tf_mode=tf_mode, name="features/stage{}/unit{}".format( i + 1, j + 1)) else: x = effi_inv_res_unit(x=x, in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, strides=strides, expansion_factor=expansion_factor, bn_epsilon=bn_epsilon, activation=activation, tf_mode=tf_mode, name="features/stage{}/unit{}".format( i + 1, j + 1)) in_channels = out_channels x = conv1x1_block(x=x, in_channels=in_channels, out_channels=final_block_channels, bn_epsilon=bn_epsilon, activation=activation, name="features/final_block") in_channels = final_block_channels x = nn.GlobalAveragePooling2D(name="features/final_pool")(x) if dropout_rate > 0.0: x = nn.Dropout(rate=dropout_rate, name="output/dropout")(x) x = nn.Dense(units=classes, input_dim=in_channels, name="output/fc", activation='softmax')(x) model = Model(inputs=input, outputs=x) model.in_size = in_size model.classes = classes return model