29 lines
852 B
Python
29 lines
852 B
Python
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'
|