DeepNet/utils/checkpoint.py at master · rvk007/DeepNet · 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
import torch
import os
import numpy as np

class Checkpoint:
    def __init__(self, path, monitor='Accuracy', mode='max', save_best_only=False, verbose=0, last_reload=False):

        self.path = path
        self.monitor = monitor
        self.mode = mode
        self.save_best_only = save_best_only
        self.verbose = verbose
        self.last_reload = last_reload

        self.last_checkpoint = None
        if self.last_reload:
            self.last_checkpoint = self.reload_checkpoint()

        self.last_metric = self.checkpoint_metric()

    def reload_checkpoint(self):
        if os.path.exists(self.path):
            last_checkpoint = torch.load(self.path,
                                    map_location=lambda storage,
                                    loc: storage)
            return last_checkpoint
        else:
            print(f'{self.path}: No such file exists.')

    def checkpoint_metric(self):
        if not os.path.exists(self.path):
            root = os.path.abspath(__file__)
            path = os.path.join(root, self.path)
            x = path.split('/')
            if x[-1].endswith('.pth'):
                os.mkdir(x[-2])
                os.mknod(os.path.join(x[-2], x[-1]))

        if self.monitor == 'Accuracy':
            metric = -np.Inf
            mode = 'max'
        elif self.monitor == 'Loss':
            metric = np.Inf
            mode = 'min'
        return metric, mode

    def save(self, model_state):
        torch.save(model_state, self.path)