CoCalc Logo Icon
StoreFeaturesDocsShareSupportNewsAboutPoliciesSign UpSign In
ai-forever

Real-time collaboration for Jupyter Notebooks, Linux Terminals, LaTeX, VS Code, R IDE, and more,
all in one place. Commercial Alternative to JupyterHub.

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