RMEP/main.py at main · cvvsu/RMEP · 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
import time, os
import torch
import numpy as np
import pandas as pd
import random
from datetime import datetime

from options import get_args
from data import build_dataset
from model import build_model


if __name__=='__main__':

    tic = time.time()
    opt, message = get_args()
    print(message)
    if not opt.isTest:
        #random.seed(opt.seed)
        #np.random.seed(opt.seed)
        #torch.manual_seed(opt.seed)
        #torch.cuda.manual_seed_all(opt.seed)
        #torch.backends.cudnn.benchmark = True

        # build the dataset according to the options
        tr_loader, val_loader = build_dataset(opt)
        model = build_model(opt)
        model.setup()

        metric = []
        best_val_loss = 1000.0
        for epoch in range(opt.start_epoch, opt.epochs+opt.epochs_decay+opt.start_epoch):
            epoch_start_time = time.time()
            iter_data_time = time.time()

            old_lr, lr = model.update_learning_rate()

            tr_loss, tr_time, tr_data_time = model.train_one_epoch(tr_loader)
            val_loss, val_time, val_data_time = model.test_one_epoch(val_loader)

            if val_loss < best_val_loss:
                best_val_loss = val_loss
                model.save_networks('best')

            if (epoch > opt.save_start_epoch) and (epoch % opt.save_freq == 0):
                model.save_networks(epoch)

            print(f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}]",
                f'Epoch [{epoch}/{opt.epochs+opt.epochs_decay}]',
                f'Time [epoch {time.time() - epoch_start_time:.6f}, data-tr {tr_data_time:.6f}, data-val {val_data_time:.6f}, train {tr_time:.6f}, val {val_time:.6f}]',
                f'Loss [train {tr_loss:.6f}, val {val_loss:.6f}]',
                f'lr [{old_lr:.6f} -> {lr:.6f}]')
            metric.append([epoch, tr_loss, val_loss, lr])
        columns = ['epoch', 'tr_loss', 'val_loss', 'lr']
        pd.DataFrame(data=metric, columns=columns).to_csv(
            os.path.join(model.save_dir, 'training.csv'), index=False)
    else:
        tr_loader, val_loader, te_loader = build_dataset(opt)
        model = build_model(opt)
        model.setup()
        tr_loss, tr_time, tr_data_time = model.test_one_epoch(tr_loader, savename='train')
        val_loss, val_time, val_data_time = model.test_one_epoch(
            val_loader, savename='val')
        te_loss, te_time, te_data_time = model.test_one_epoch(
            te_loader, savename='test')
        print(f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}]",
              f"Time {te_time+val_time+te_time}",
              f"loss [train: {tr_loss:.6f}, val: {val_loss:6f}, test: {te_loss:.6f}]")
    print(f'Time consuming for all epochs: {(time.time() - tic)/60:.2f} minutes.')