config.py 852 B

12345678910111213141516171819202122232425262728
  1. import os
  2. import torch
  3. device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # sets device for model and PyTorch tensors
  4. # Model parameters
  5. image_w = 112
  6. image_h = 112
  7. channel = 3
  8. emb_size = 512
  9. # Training parameters
  10. num_workers = 1 # for data-loading; right now, only 1 works with h5py
  11. grad_clip = 5. # clip gradients at an absolute value of
  12. print_freq = 100 # print training/validation stats every __ batches
  13. checkpoint = None # path to checkpoint, None if none
  14. # Data parameters
  15. num_classes = 93431
  16. num_samples = 5179510
  17. DATA_DIR = 'data'
  18. # faces_ms1m_folder = 'data/faces_ms1m_112x112'
  19. faces_ms1m_folder = 'data/ms1m-retinaface-t1'
  20. path_imgidx = os.path.join(faces_ms1m_folder, 'train.idx')
  21. path_imgrec = os.path.join(faces_ms1m_folder, 'train.rec')
  22. IMG_DIR = 'data/images'
  23. pickle_file = 'data/faces_ms1m_112x112.pickle'