Rednet_Image_Enhancer/train.py at main · dungrup/Rednet_Image_Enhancer · 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
import argparse
import os
import torch
from torch import nn
import torch.optim as optim
import torch.backends.cudnn as cudnn
from torch.utils.data.dataloader import DataLoader
from tqdm import tqdm
from model import REDNet10, REDNet20, REDNet30
from dataset import BDDDataset
from utils import AverageMeter
from torch.utils.tensorboard import SummaryWriter

cudnn.benchmark = True
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--arch', type=str, default='REDNet10', help='REDNet10, REDNet20, REDNet30')
    parser.add_argument('--raw_images_dir', type=str, required=True)
    parser.add_argument('--comp_images_dir', type=str, required=True)
    parser.add_argument('--outputs_dir', type=str, required=True)
    parser.add_argument('--patch_size', type=int, default=50)
    parser.add_argument('--batch_size', type=int, default=2)
    parser.add_argument('--num_epochs', type=int, default=20)
    parser.add_argument('--lr', type=float, default=1e-4)
    parser.add_argument('--threads', type=int, default=8)
    parser.add_argument('--seed', type=int, default=42)
    opt = parser.parse_args()

    if not os.path.exists(opt.outputs_dir):
        os.makedirs(opt.outputs_dir)

    torch.manual_seed(opt.seed)

    if opt.arch == 'REDNet10':
        model = REDNet10()
    elif opt.arch == 'REDNet20':
        model = REDNet20()
    elif opt.arch == 'REDNet30':
        model = REDNet30()

    # If there is an existing model, load it and train it
    if os.path.exists(os.path.join(opt.outputs_dir, f'{opt.arch}_best.pth')):
        model.load_state_dict(torch.load(os.path.join(opt.outputs_dir, f'{opt.arch}_best.pth')))

    model = model.to(device)
    criterion = nn.MSELoss()

    optimizer = optim.Adam(model.parameters(), lr=opt.lr)

    dataset = BDDDataset(opt.raw_images_dir, opt.comp_images_dir, opt.patch_size)
    val_dataset = BDDDataset(opt.raw_images_dir, opt.comp_images_dir, opt.patch_size, train_flag=False)
    dataloader = DataLoader(dataset=dataset,
                            batch_size=opt.batch_size,
                            shuffle=True,
                            num_workers=opt.threads,
                            pin_memory=True,
                            drop_last=True)

    val_dataloader = DataLoader(dataset=val_dataset,
                            batch_size=1,
                            shuffle=False,
                            num_workers=opt.threads,
                            pin_memory=True,
                            drop_last=True)

    best_val_loss = float('inf')

    # Initialize Tensorboard writer
    writer = SummaryWriter(log_dir=os.path.join('./', 'logs'))

    for epoch in range(opt.num_epochs):
        epoch_losses = AverageMeter()

        with tqdm(total=(len(dataset) - len(dataset) % opt.batch_size)) as _tqdm:
            _tqdm.set_description('epoch: {}/{}'.format(epoch + 1, opt.num_epochs))
            for data in dataloader:
                inputs, labels = data
                inputs = inputs.to(device)
                labels = labels.to(device)

                preds = model(inputs)

                loss = criterion(preds, labels)
                epoch_losses.update(loss.item(), len(inputs))

                optimizer.zero_grad()
                loss.backward()
                optimizer.step()

                _tqdm.set_postfix(loss='{:.6f}'.format(epoch_losses.avg))
                _tqdm.update(len(inputs))


        # Log training loss to Tensorboard
        writer.add_scalar('Loss/train', epoch_losses.avg, epoch)

        # Validation phase
        model.eval()
        val_losses = AverageMeter()

        with torch.no_grad():
            for data in val_dataloader:
                inputs, labels = data
                inputs = inputs.to(device)
                labels = labels.to(device)

                preds = model(inputs)
                loss = criterion(preds, labels)

                val_losses.update(loss.item(), len(inputs))

        print(f'Validation loss: {val_losses.avg:.6f}')

        # Log validation loss to Tensorboard
        writer.add_scalar('Loss/val', val_losses.avg, epoch)

        # Save model if validation loss improves
        if val_losses.avg < best_val_loss:
            best_val_loss = val_losses.avg
            torch.save(model.state_dict(),
                      os.path.join(opt.outputs_dir, f'{opt.arch}_best.pth'))

        torch.save(model.state_dict(), os.path.join(opt.outputs_dir, '{}_epoch_{}.pth'.format(opt.arch, epoch)))

writer.close()