-
Notifications
You must be signed in to change notification settings - Fork 2
/
main.py
45 lines (32 loc) · 1.12 KB
/
main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
import torch
from pathlib import Path
from utils.data import make_blobs_dataset
from utils.nnet import get_device
from hebbcl.logger import LoggerFactory
from hebbcl.model import ModelFactory
from hebbcl.trainer import Optimiser, train_on_blobs
from hebbcl.parameters import parser
args = parser.parse_args()
# overwrite cuda argument depending on GPU availability
args.cuda = args.cuda and torch.cuda.is_available()
if __name__ == "__main__":
# create checkpoint dir
save_dir = (
Path("checkpoints") / "quick_test"
) # datetime.datetime.now().strftime("%Y-%m-%dT%H-%M-%S")
# get (cuda) device
args.device, _ = get_device(args.cuda)
# get dataset
dataset = make_blobs_dataset(args)
# instantiate logger, model and optimiser
logger = LoggerFactory.create(args, save_dir)
model = ModelFactory.create(args)
optim = Optimiser(args)
# send model to GPU
model = model.to(args.device)
# train model
train_on_blobs(args, model, optim, dataset, logger)
# save results
if args.save_results:
save_dir.mkdir(parents=True, exist_ok=True)
logger.save(model)