update

models/config.py

@@ -0,0 +1,28 @@
+import os
+
+import torch
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')  # sets device for model and PyTorch tensors
+
+# Model parameters
+image_w = 112
+image_h = 112
+channel = 3
+emb_size = 512
+
+# Training parameters
+num_workers = 1  # for data-loading; right now, only 1 works with h5py
+grad_clip = 5.  # clip gradients at an absolute value of
+print_freq = 100  # print training/validation stats  every __ batches
+checkpoint = None  # path to checkpoint, None if none
+
+# Data parameters
+num_classes = 93431
+num_samples = 5179510
+DATA_DIR = 'data'
+# faces_ms1m_folder = 'data/faces_ms1m_112x112'
+faces_ms1m_folder = 'data/ms1m-retinaface-t1'
+path_imgidx = os.path.join(faces_ms1m_folder, 'train.idx')
+path_imgrec = os.path.join(faces_ms1m_folder, 'train.rec')
+IMG_DIR = 'data/images'
+pickle_file = 'data/faces_ms1m_112x112.pickle'