Capestone/main.py at main · rishu072/Capestone · GitHub

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
from src.data import load_data
from src.preprocess import get_loader
from src.autoencoder import Autoencoder
from src.train_autoencoder import train_autoencoder
from src.classifier import Classifier
from src.train_classifier import extract_features, train_classifier
from src.evaluate import evaluate
from src.visualize import plot_loss, plot_latent, plot_confusion
import torch
import config

# load data
train, test = load_data()
train_loader, test_loader = get_loader(train, test)

# train autoencoder
ae = Autoencoder()
ae, ae_loss = train_autoencoder(ae, train_loader)

# save AE
torch.save(ae.state_dict(), config.MODEL_PATH + "autoencoder.pt")

# extract features
train_f, train_y = extract_features(ae, train_loader)
test_f, test_y = extract_features(ae, test_loader)

# train classifier
clf = Classifier()
clf, clf_loss = train_classifier(clf, train_f, train_y)

# save classifier
torch.save(clf.state_dict(), config.MODEL_PATH + "classifier.pt")

# evaluate
preds = evaluate(clf, test_f, test_y)

# visualize
plot_loss(ae_loss, "AE Loss")
plot_loss(clf_loss, "Classifier Loss")
plot_latent(test_f, test_y)
plot_confusion(test_y, preds)

from src.raw_classifier import RawClassifier
from src.train_classifier import train_raw

# RAW classifier
raw_model = RawClassifier()
raw_model = train_raw(raw_model, train_loader)

# evaluate raw
from src.evaluate import evaluate
print("\nRAW MODEL:")
evaluate(raw_model, *next(iter(test_loader)))

from src.visualize import show_reconstruction
show_reconstruction(ae, test_loader)