-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.py
60 lines (41 loc) · 1.58 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
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
import random
import torch
import numpy as np
import torch.nn as nn
import torch.optim as optim
from comet_ml import Experiment
from train import training
from utils import load_model
from data import load_dataset
from delta import setup_delta_tracking
def run_experiment(epochs, model_name, training_type, configs):
""" Runs the basic experiment"""
print(epochs, "CONFIGS: ", configs)
experiment = Experiment(api_key=configs.api_key,
project_name=configs.project_name, workspace="ayushm-agrawal")
experiment.set_name(configs.exp_name)
experiment.log_parameters(configs)
# set seed for reproducibility.
seed = configs.seed
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
# gpu training specific seed settings.
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
# load data
loaders = load_dataset(configs)
# load model
model = load_model(model_name, training_type, configs)
# loss
criterion = nn.CrossEntropyLoss().cuda()
# optimizer
optimizer = optim.SGD(model.parameters(), configs.lr,
momentum=configs.momentum,
weight_decay=configs.weight_decay)
# get tracking dictionaries
prev_list, rmae_delta_dict = setup_delta_tracking(model)
# train model
rmae_delta_dict, train_acc_arr, test_acc_arr = training(epochs, loaders, model, optimizer, criterion, prev_list, rmae_delta_dict, configs, experiment)
experiment.end()
return rmae_delta_dict, train_acc_arr, test_acc_arr