-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.py
117 lines (104 loc) · 4.57 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
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
"""
Main calling functions for training and evaluating the ROAD models
"""
from torch.utils.data import DataLoader
import numpy as np
from utils.args import args
from utils.data import defaults
from data import get_data
from models import (VAE, BackBone,
PositionClassifier, Decoder,
ClassificationHead)
from train import train_vae, train_supervised, train_ssl
from fine_tune import fine_tune
from eval import (eval_supervised,
eval_classification_head, eval_knn)
print(args.model_name)
(train_dataset,
val_dataset,
test_dataset,
supervised_train_dataset,
supervised_val_dataset) = get_data(args)
train_dataloader = DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=True)
supervised_train_dataloader = DataLoader(supervised_train_dataset,
batch_size=args.batch_size,
shuffle=True)
test_dataloader = DataLoader(test_dataset,
batch_size=args.batch_size,
shuffle=False)
if args.model == 'vae':
vae = VAE(in_channels=4,
latent_dim=args.latent_dim,
patch_size=args.patch_size,
hidden_dims=args.hidden_dims)
if args.load_model:
vae.load(args)
else:
vae = train_vae(train_dataloader, vae, args)
pred, thr = eval_knn(vae, None, test_dataloader, train_dataloader, args)
elif args.model in ('supervised', 'all', 'random_init', 'dknn'):
supervised_backbone = BackBone(in_channels=4,
out_dims=len(defaults.anomalies)+1,
model_type=args.backbone,
supervision=True)
if args.load_model:
supervised_backbone.load(args, 'supervised', False)
else:
supervised_backbone = train_supervised(supervised_train_dataloader,
supervised_val_dataset,
supervised_backbone,
args)
if args.model in ('ssl', 'all', 'random_init', 'dknn'):
if args.model == 'dknn':
weight_init = 'DEFAULT'
else:
weight_init = None
ssl_backbone = BackBone(in_channels=4,
out_dims=args.latent_dim,
model_type=args.backbone,
weight_init=weight_init)
position_classifier = PositionClassifier(latent_dim=args.latent_dim,
out_dims=8)
decoder = Decoder(out_channels=4,
patch_size=args.patch_size,
latent_dim=args.latent_dim,
n_layers=5)
classification_head = ClassificationHead(out_dims=1,
latent_dim=args.latent_dim)
if args.load_model:
ssl_backbone.load(args, 'ssl', True)
position_classifier.load(args)
decoder.load(args)
classification_head.load(args)
else:
if args.model != 'random_init' and args.model != 'dknn':
(ssl_backbone,
position_classifier,
decoder) = train_ssl(train_dataloader,
val_dataset,
ssl_backbone,
position_classifier,
decoder,
args)
(ssl_backbone,
classification_head) = fine_tune(supervised_train_dataloader,
supervised_val_dataset,
test_dataloader,
ssl_backbone,
classification_head,
args)
if args.model in ('all', 'random_init', 'dknn'):
for i in range(10):
test_dataloader.dataset.set_seed(np.random.randint(1000))
pred, thr = eval_supervised(supervised_backbone, test_dataloader, args)
pred_ft, thr_ft = eval_classification_head(ssl_backbone,
classification_head,
test_dataloader,
args)
pred, thr = eval_supervised(supervised_backbone,
test_dataloader,
args,
pred_ft,
thr_ft)