Skip to content
Snippets Groups Projects
Select Git revision
  • main default protected
1 result
Blame Permalink
config.py 749 B
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17 






import torch 
import torchvision.transforms as transforms


DEVICE = "cuda" if torch.cuda.is_available() else "cpu"  # Device 
NB_EPOCHS = 50 # Number of epochs for training 
BATCH_SIZE = 20  # Size of the training batch
TRAINING_IMAGE_FILE_PATH = "./data/train/images/"
TRAINING_LABEL_FILE_PATH = "./data/train/labels/"
VAL_IMAGE_FILE_PATH = "./data/val/images/"
VAL_LABEL_FILE_PATH = "./data/val/labels/"

TRANFORMATIONS = transforms.Compose([transforms.Resize((32,32)),
                                     transforms.ToTensor(),
                                     transforms.Normalize(mean=[0.4914, 0.4822, 0.4465],
                                                          std=[0.2023, 0.1994, 0.2010])
                                     ])