-
Notifications
You must be signed in to change notification settings - Fork 0
/
train.py
47 lines (31 loc) · 1.43 KB
/
train.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
import torch
from torchmetrics.functional.classification import multiclass_f1_score, multiclass_average_precision
from transformers.modeling_outputs import ImageClassifierOutput
import wandb
def train(model, device, dataloader, number_of_classes, criterion, optimizer):
model.train()
running_loss = 0.0
predictions, targets = [], []
for i, (inputs, labels) in enumerate(dataloader):
inputs = inputs.to(device)
labels = labels.to(device)
optimizer.zero_grad()
outputs = model(inputs)
#transformer model outputs an ImageClassifierOutput object
if isinstance(outputs, ImageClassifierOutput):
outputs = outputs.logits
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item() * inputs.size(0)
predictions.append(outputs.cpu())
targets.append(labels.cpu())
predictions = torch.cat(predictions, 0)
targets = torch.cat(targets, 0)
# Calculate epoch loss and metrics
epoch_f1 = multiclass_f1_score(predictions, targets, num_classes=number_of_classes)
epoch_auprc = multiclass_average_precision(predictions, targets, num_classes=number_of_classes)
epoch_loss = running_loss / len(dataloader.dataset)
if wandb.run is not None:
wandb.log({'train_loss': epoch_loss, 'train_f1': epoch_f1, 'train_auprc': epoch_auprc})
return epoch_loss, epoch_f1, epoch_auprc