-
Notifications
You must be signed in to change notification settings - Fork 10
/
Copy path08_baseline.py
134 lines (99 loc) · 4.4 KB
/
08_baseline.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
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
import time
import lightning as L
from lightning import Fabric
import torch
import torch.nn.functional as F
import torchmetrics
from torchvision import transforms
from torchvision.models import vit_h_14
from torchvision.models import ViT_H_14_Weights
from watermark import watermark
from local_utilities import get_dataloaders_cifar10
def train(num_epochs, model, optimizer, train_loader, val_loader, fabric):
for epoch in range(num_epochs):
train_acc = torchmetrics.Accuracy(task="multiclass", num_classes=10).to(fabric.device)
model.train()
for batch_idx, (features, targets) in enumerate(train_loader):
model.train()
### FORWARD AND BACK PROP
logits = model(features)
loss = F.cross_entropy(logits, targets)
optimizer.zero_grad()
fabric.backward(loss)
### UPDATE MODEL PARAMETERS
optimizer.step()
### LOGGING
if not batch_idx % 300:
print(f"Epoch: {epoch+1:04d}/{num_epochs:04d} | Batch {batch_idx:04d}/{len(train_loader):04d} | Loss: {loss:.4f}")
model.eval()
with torch.no_grad():
predicted_labels = torch.argmax(logits, 1)
train_acc.update(predicted_labels, targets)
### MORE LOGGING
model.eval()
with torch.no_grad():
val_acc = torchmetrics.Accuracy(task="multiclass", num_classes=10).to(fabric.device)
for (features, targets) in val_loader:
outputs = model(features)
predicted_labels = torch.argmax(outputs, 1)
val_acc.update(predicted_labels, targets)
print(f"Epoch: {epoch+1:04d}/{num_epochs:04d} | Train acc.: {train_acc.compute()*100:.2f}% | Val acc.: {val_acc.compute()*100:.2f}%")
train_acc.reset(), val_acc.reset()
if __name__ == "__main__":
print(watermark(packages="torch,lightning", python=True))
print("Torch CUDA available?", torch.cuda.is_available())
fabric = Fabric(accelerator="cuda", devices=1, precision="16-mixed")
fabric.launch()
L.seed_everything(123)
##########################
### 1 Loading the Dataset
##########################
train_transforms = transforms.Compose([transforms.Resize((518, 518)),
#transforms.RandomCrop((224, 224)),
transforms.ToTensor()])
test_transforms = transforms.Compose([transforms.Resize((518, 518)),
#transforms.CenterCrop((224, 224)),
transforms.ToTensor()])
train_loader, val_loader, test_loader = get_dataloaders_cifar10(
batch_size=4,
num_workers=1,
train_transforms=train_transforms,
test_transforms=test_transforms,
validation_fraction=0.1)
train_loader, val_loader, test_loader = fabric.setup_dataloaders(
train_loader, val_loader, test_loader)
#########################################
### 2 Initializing the Model
#########################################
model = vit_h_14(weights=ViT_H_14_Weights.IMAGENET1K_SWAG_E2E_V1)
# replace output layer
model.heads.head = torch.nn.Linear(in_features=1280, out_features=10)
optimizer = torch.optim.Adam(model.parameters(), lr=5e-5)
model, optimizer = fabric.setup(model, optimizer)
#########################################
### 3 Finetuning
#########################################
start = time.time()
train(
num_epochs=1,
model=model,
optimizer=optimizer,
train_loader=train_loader,
val_loader=val_loader,
fabric=fabric
)
end = time.time()
elapsed = end-start
print(f"Time elapsed {elapsed/60:.2f} min")
print(f"Memory used: {torch.cuda.max_memory_reserved() / 1e9:.02f} GB")
#########################################
### 4 Evaluation
#########################################
with torch.no_grad():
model.eval()
test_acc = torchmetrics.Accuracy(task="multiclass", num_classes=10).to(fabric.device)
for (features, targets) in test_loader:
outputs = model(features)
predicted_labels = torch.argmax(outputs, 1)
test_acc.update(predicted_labels, targets)
print(f"Test accuracy {test_acc.compute()*100:.2f}%")