File size: 4,678 Bytes
b2ffc9b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
135
136
137
138
139
140
141
142
143
144
145
146
import time

import numpy as np
import torch
from torch.nn import functional as F


def train_epoch(train_loader, model, loss_function, optimizer, device, epoch):
    model.train()

    correct = 0
    total = 0
    losses = 0
    t0 = time.time()
    for idx, (batch_images, batch_labels) in enumerate(train_loader):
        # Loading tensors in the used device
        step_images, step_labels = batch_images.to(device), batch_labels.to(device)

        # zero the parameter gradients
        optimizer.zero_grad()

        step_output = model(step_images)
        loss = loss_function(step_output, step_labels)
        loss.backward()
        optimizer.step()

        step_total = step_labels.size(0)
        step_loss = loss.item()
        losses += step_loss*step_total
        total += step_total

        step_preds = torch.max(step_output.data, 1)[1]
        step_correct = (step_preds == step_labels).sum().item()
        correct += step_correct

    train_loss = losses / total
    train_acc = correct / total
    format_args = (epoch, train_acc, train_loss, time.time() - t0)
    print('EPOCH {} :: train accuracy: {:.4f} - train loss: {:.4f} at {:.1f}s'.format(*format_args))


def val_epoch(val_loader, model, loss_function, device, epoch):
    model.eval()

    y_true = []
    y_pred = []

    correct = 0
    total = 0
    losses = 0
    t0 = time.time()
    with torch.no_grad():
        for batch_images, batch_labels in val_loader:
            # Loading tensors in the used device
            step_images, step_labels = batch_images.to(device), batch_labels.to(device)

            step_output = model(step_images)
            loss = loss_function(step_output, step_labels)

            step_total = step_labels.size(0)
            step_loss = loss.item()
            losses += step_loss*step_total
            total += step_total

            step_preds = torch.max(step_output.data, 1)[1]
            y_pred.append(step_preds.cpu().detach().numpy())
            y_true.append(step_labels.cpu().detach().numpy())
            step_correct = (step_preds == step_labels).sum().item()
            correct += step_correct

    val_loss = losses / total
    val_acc = correct / total
    format_args = (epoch, val_acc, val_loss, time.time() - t0)
    print('EPOCH {} :: val accuracy: {:.4f} - val loss: {:.4f} at {:.1f}s'.format(*format_args))

    y_pred = np.concatenate(y_pred, axis=0)
    y_true = np.concatenate(y_true, axis=0)
    return y_true, y_pred


def test_epoch(test_loader, model, loss_function, device):
    model.eval()

    correct = 0
    total = 0
    losses = 0
    all_true = []
    all_pred = []
    t0 = time.time()
    with torch.no_grad():
        for batch_images, batch_labels in test_loader:
            # Loading tensors in the used device
            step_images, step_labels = batch_images.to(device), batch_labels.to(device)

            step_output = model(step_images)
            loss = loss_function(step_output, step_labels)

            step_total = step_labels.size(0)
            step_loss = loss.item()
            losses += step_loss*step_total
            total += step_total

            step_preds = torch.max(step_output.data, 1)[1]
            step_correct = (step_preds == step_labels).sum().item()
            correct += step_correct

            all_true.append(step_labels.cpu().numpy())
            all_pred.append(step_preds.cpu().numpy())

    val_loss = losses / total
    val_acc = correct / total
    format_args = (val_acc, val_loss, time.time() - t0)
    print('EPOCH :: test accuracy: {:.4f} - test loss: {:.4f} at {:.1f}s'.format(*format_args))

    all_pred = np.concatenate(all_pred, axis=0)
    all_true = np.concatenate(all_true, axis=0)
    return all_true, all_pred


def detection_epoch(detection_loader, model, device):
    model.eval()

    pred_probs = []
    coords_x = []
    coords_y = []
    t0 = time.time()
    with torch.no_grad():
        for batch_images, batch_x, batch_y in detection_loader:
            # Loading tensors in the used device
            step_images = batch_images.to(device)
            step_output = model(step_images)
            step_pred_probs = F.softmax(step_output, 1)

            step_pred_probs = step_pred_probs.cpu().numpy()
            step_x = batch_x.numpy()
            step_y = batch_y.numpy()

            coords_x.append(step_x)
            coords_y.append(step_y)
            pred_probs.append(step_pred_probs)

    return_pred_probs = np.concatenate(pred_probs, axis=0)
    return_coords_x = np.concatenate(coords_x, axis=0)
    return_coords_y = np.concatenate(coords_y, axis=0)
    return return_pred_probs, return_coords_x, return_coords_y