weakly-supervised-metric-embedding/training.py at main · arkwave/weakly-supervised-metric-embedding · 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
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
# -*- coding: utf-8 -*-
# @Author: Ananth
# @Date:   2020-08-25 11:25:26
# @Last Modified by:   Ananth
# @Last Modified time: 2020-10-15 01:26:10

from utils import jacobian
import numpy as np
import torch
import torch.nn as nn
from torchvision.datasets import MNIST
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt


def train_network(network, datagen, epochs, metric_loss, reconstruction_loss,
                  optimizer, mu=1, contractive=False, device=None):
    """[summary]

    Args:
        network ([type]): [description]
        datagen ([type]): [description]
        epochs ([type]): [description]
        metric_loss ([type]): [description]
        reconstruction_loss ([type]): [description]
        optimizer ([type]): [description]
        mu (int, optional): [description]. Defaults to 1.
        contractive (bool, optional): [description]. Defaults to False.
        device ([type], optional): [description]. Defaults to None.

    Returns:
        [type]: [description]
    """

    recon_losses = []
    metric_losses = []
    total_losses = []

    print("using %s" % device)

    network = network.to(device)
    network.train()

    for epoch in range(epochs):

        # generate the batch
        batch, labels = datagen.generate_batch()
        batch = tuple(x.to(device) for x in batch)
        labels = labels.to(device)

        # zero gradients
        optimizer.zero_grad()

        # forward step
        outputs = network(*batch)

        # separate into encoded and reconstructed
        encoded, reconstructed = list(zip(*outputs))
        # encoded = [outputs[0] for output in outputs]
        # reconstructed = [i[1] for i in outputs]
        m_loss = metric_loss(*encoded, labels) * mu

        j_norm = 0

        if contractive:
            for i in range(len(encoded)):
                J = jacobian(batch[i], encoded[i])
                j_norm += torch.norm(J)
            print("j_norm: ", j_norm)

        recon_loss = reconstruction_loss(batch, reconstructed, additional_losses=j_norm) * (1-mu)

        # add the losses together
        loss = recon_loss + m_loss

        # append losses to history.
        total_losses.append(loss.item())
        recon_losses.append(recon_loss.item())
        metric_losses.append(m_loss.item())

        # print("Iteration %s. Metric Loss: %s | Reconstruction Loss: %s" % (epoch, round(m_loss.item(), 3), round(recon_loss.item(), 3)))
        if epoch % (epochs/100) == 0 and epoch != 0:
            print("Iteration %s. Metric Loss: %s | Reconstruction Loss: %s" % (epoch, round(m_loss.item(), 3), round(recon_loss.item(), 3)))

        # backward step
        loss.backward()
        optimizer.step()

    return network, total_losses, recon_losses, metric_losses


def train_vae(network, train_set, test_set, optimizer, loss, epochs, batch_size=64):
    """Training/evaluation loop for the variational autoencoder.

    Args:
        network ([type]): [description]
        trainloader ([type]): [description]
        testloader ([type]): [description]
        optimizer ([type]): [description]
        epochs ([type]): [description]
    """
    recon_loss = nn.MSELoss()
    losses = []
    val_losses = []

    # instantiate dataloaders.
    trainloader = DataLoader(train_set, batch_size=batch_size, shuffle=True)
    testloader = DataLoader(test_set, batch_size=1, shuffle=False)

    for epoch in range(epochs):
        for idx, (sample, label) in enumerate(trainloader):
            optimizer.zero_grad()
            total_dim = np.product(sample.shape[1:])
            sample = sample.reshape(-1, total_dim)

            # forward pass
            (mu, logvar), reconstructed = network(sample)
            # compute loss
            total_loss = loss(sample, reconstructed, mu, logvar, batch_size, total_dim)
            # backward pass.
            total_loss.backward()
            optimizer.step()
            losses.append(total_loss.item())

            if idx % 100 == 0 and idx > 0:
                print("epoch: %s | idx: %s | loss: %s" % (epoch, idx, total_loss.item()))

    return network, losses