-
Notifications
You must be signed in to change notification settings - Fork 4
/
train.py
280 lines (231 loc) · 10.1 KB
/
train.py
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
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
import os
import torch
import yaml
from utils import network_parameters, losses
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
import time
import numpy as np
import random
from transform.data_RGB import get_training_data,get_validation_data2
from warmup_scheduler.scheduler import GradualWarmupScheduler
from tqdm import tqdm
from pytorch_msssim import ssim
from tensorboardX import SummaryWriter
import utils.losses
from model.Walmafa import Walmafa
import argparse
parser = argparse.ArgumentParser(description='Hyper-parameters for LLFormer')
parser.add_argument('-yml_path', default="./configs/LOL/train/training_LOL.yaml", type=str)
args = parser.parse_args()
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
class TVLoss(nn.Module):
def __init__(self, TVLoss_weight=1):
super(TVLoss, self).__init__()
self.TVLoss_weight = TVLoss_weight
def forward(self, x):
batch_size = x.size()[0]
h_x = x.size()[2]
w_x = x.size()[3]
count_h = self._tensor_size(x[:, :, 1:, :])
count_w = self._tensor_size(x[:, :, :, 1:])
# 横向上的相邻像素的差异值
h_tv = torch.pow((x[:, :, 1:, :] - x[:, :, :h_x - 1, :]), 2).sum()
##纵向上的相邻像素的差异值
w_tv = torch.pow((x[:, :, :, 1:] - x[:, :, :, :w_x - 1]), 2).sum()
return self.TVLoss_weight * 2 * (h_tv / count_h + w_tv / count_w) / batch_size
def _tensor_size(self, t):
return t.size()[1] * t.size()[2] * t.size()[3]
## Set Seeds
torch.backends.cudnn.benchmark = True
random.seed(1234)
np.random.seed(1234)
torch.manual_seed(1234)
torch.cuda.manual_seed_all(1234)
## Load yaml configuration file
yaml_file = args.yml_path
with open(yaml_file, 'r') as config:
opt = yaml.safe_load(config)
print("load training yaml file: %s"%(yaml_file))
Train = opt['TRAINING']
OPT = opt['OPTIM']
## Build Model
print('==> Build the model')
model_restored = Walmafa(inp_channels=3,out_channels=3,dim = 16,num_blocks = [2,3,4],heads = [1,2,4,8],ffn_expansion_factor = 2.66,bias = False,LayerNorm_type = 'WithBias',attention=True,skip = False)
p_number = network_parameters(model_restored)
model_restored.cuda()
## Training model path direction
mode = opt['MODEL']['MODE']
model_dir = os.path.join(Train['SAVE_DIR'], mode, 'models')
utils.mkdir(model_dir)
train_dir = Train['TRAIN_DIR']
val_dir = Train['VAL_DIR']
## GPU
gpus = ','.join([str(i) for i in opt['GPU']])
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = gpus
device_ids = [i for i in range(torch.cuda.device_count())]
if torch.cuda.device_count() > 1:
print("\n\nLet's use", torch.cuda.device_count(), "GPUs!\n\n")
if len(device_ids) > 1:
model_restored = nn.DataParallel(model_restored, device_ids=device_ids)
## Optimizer
start_epoch = 1
new_lr = float(OPT['LR_INITIAL'])
optimizer = optim.Adam(model_restored.parameters(), lr=new_lr, betas=(0.9, 0.999), eps=5e-5)
## Scheduler (Strategy)
warmup_epochs = 1
scheduler_cosine = optim.lr_scheduler.CosineAnnealingLR(optimizer, OPT['EPOCHS'] - warmup_epochs,
eta_min=float(OPT['LR_MIN']))
scheduler = GradualWarmupScheduler(optimizer, multiplier=1, total_epoch=warmup_epochs, after_scheduler=scheduler_cosine)
scheduler.step()
## Resume (Continue training by a pretrained model)
if Train['RESUME']:
path_chk_rest = utils.get_last_path(model_dir, '_latest.pth')
path_best_PSNR = model_dir+'/model_bestPSNR.pth'
path_best_SSIM = model_dir + '/model_bestSSIM.pth'
best_epoch_psnr, best_psnr = utils.load_best_PSNR_epoch_PSNR(path_best_PSNR)
best_epoch_ssim, best_ssim = utils.load_best_SSIM_epoch_SSIM(path_best_SSIM)
utils.load_checkpoint(model_restored, path_chk_rest)
start_epoch = utils.load_start_epoch(path_chk_rest) + 1
end_epoch = start_epoch+OPT['EPOCHS']-1
utils.load_optim(optimizer, path_chk_rest)
for i in range(1, start_epoch):
scheduler.step()
new_lr = scheduler.get_lr()[0]
print('------------------------------------------------------------------')
print("==> Resuming Training with learning rate:", new_lr)
print('------------------------------------------------------------------')
else:
best_psnr = 0
best_ssim = 0
best_epoch_psnr = 0
best_epoch_ssim = 0
end_epoch = OPT['EPOCHS']
## Loss
L1loss = nn.L1Loss()
Charloss = nn.SmoothL1Loss()
Ms_ssim_L1_loss = utils.losses.MS_SSIM_L1_LOSS
## DataLoaders
print('==> Loading datasets')
train_dataset = get_training_data(train_dir, {'patch_size': Train['TRAIN_PS']})
train_loader = DataLoader(dataset=train_dataset, batch_size=OPT['BATCH'],
shuffle=True, num_workers=8, drop_last=False)
val_dataset = get_validation_data2(val_dir, {'patch_size': Train['VAL_PS']})
val_loader = DataLoader(dataset=val_dataset, batch_size=1, shuffle=False, num_workers=8,
drop_last=False)
print('------------------------------------------------------------------')
# Show the training configuration
print(f'''==> Training details:
------------------------------------------------------------------
Restoration mode: {mode}
Train patches size: {str(Train['TRAIN_PS']) + 'x' + str(Train['TRAIN_PS'])}
Val patches size: {str(Train['VAL_PS']) + 'x' + str(Train['VAL_PS'])}
Model parameters: {p_number/(1024*1024):.2f}M
Start/End epochs: {str(start_epoch) + '~' + str(end_epoch)}
Best_PSNR_Epoch: {best_epoch_psnr}
Best_PSNR: {best_psnr:.4f}
Best_SSIM_Epoch: {best_epoch_ssim}
Best_SSIM: {best_ssim:.4f}
Batch sizes: {OPT['BATCH']}
Learning rate: {OPT['LR_INITIAL']}
GPU: {'GPU' + str(device_ids)}''')
print('------------------------------------------------------------------')
# Start training!
print('==> Training start: ')
# best_psnr = 0
# best_ssim = 0
# best_epoch_psnr = 0
# best_epoch_ssim = 0
total_start_time = time.time()
## Log
log_dir = os.path.join(Train['SAVE_DIR'], mode, 'log')
utils.mkdir(log_dir)
writer = SummaryWriter(log_dir=log_dir, filename_suffix=f'_{mode}')
for epoch in range(start_epoch, end_epoch + 1):
epoch_start_time = time.time()
epoch_loss = 0
train_id = 1
model_restored.train()
for i, data in enumerate(tqdm(train_loader), 0):
# Forward propagation
for param in model_restored.parameters():
param.grad = None
target = data[0].cuda()
input_ = data[1].cuda()
restored = model_restored(input_)
# Compute loss
content_loss = L1loss(restored, target)
ssim_loss = 1 - ssim(restored, target, data_range=1.0).to("cuda:0")
content_loss = content_loss+ssim_loss
tv_loss = TVLoss()
TV_loss = tv_loss(restored)
loss = 0.95*content_loss + 0.05*TV_loss
# Back propagation
loss.backward()
optimizer.step()
epoch_loss += loss.item()
## Evaluation (Validation)
if epoch % Train['VAL_AFTER_EVERY'] == 0:
model_restored.eval()
psnr_val_rgb = []
ssim_val_rgb = []
for ii, data_val in enumerate(val_loader, 0):
target = data_val[0].cuda()
input_ = data_val[1].cuda()
h, w = target.shape[2], target.shape[3]
with torch.no_grad():
restored = model_restored(input_)
restored = restored[:, :, :h, :w]
for res, tar in zip(restored, target):
psnr_val_rgb.append(utils.torchPSNR(res, tar))
ssim_val_rgb.append(utils.torchSSIM(restored, target))
psnr_val_rgb = torch.stack(psnr_val_rgb).mean().item()
ssim_val_rgb = torch.stack(ssim_val_rgb).mean().item()
# Save the best PSNR model of validation
if psnr_val_rgb > best_psnr:
best_psnr = psnr_val_rgb
best_epoch_psnr = epoch
torch.save({'epoch': epoch,
'PSNR': best_psnr,
'state_dict': model_restored.state_dict(),
'optimizer': optimizer.state_dict()
}, os.path.join(model_dir, "model_bestPSNR.pth"))
print("[epoch %d PSNR: %.4f --- best_epoch %d Best_PSNR %.4f]" % (
epoch, psnr_val_rgb, best_epoch_psnr, best_psnr))
# Save the best SSIM model of validation
if ssim_val_rgb > best_ssim:
best_ssim = ssim_val_rgb
best_epoch_ssim = epoch
torch.save({'epoch': epoch,
'SSIM' : best_ssim,
'state_dict': model_restored.state_dict(),
'optimizer': optimizer.state_dict()
}, os.path.join(model_dir, "model_bestSSIM.pth"))
print("[epoch %d SSIM: %.4f --- best_epoch %d Best_SSIM %.4f]" % (
epoch, ssim_val_rgb, best_epoch_ssim, best_ssim))
"""
# Save evey epochs of model
torch.save({'epoch': epoch,
'state_dict': model_restored.state_dict(),
'optimizer': optimizer.state_dict()
}, os.path.join(model_dir, f"model_epoch_{epoch}.pth"))
"""
writer.add_scalar('val/PSNR', psnr_val_rgb, epoch)
writer.add_scalar('val/SSIM', ssim_val_rgb, epoch)
scheduler.step()
print("------------------------------------------------------------------")
print("Epoch: {}\tTime: {:.4f}\tLoss: {:.4f}\tLearningRate {:.6f}".format(epoch, time.time() - epoch_start_time,
epoch_loss, scheduler.get_lr()[0]))
print("------------------------------------------------------------------")
# Save the last model
torch.save({'epoch': epoch,
'state_dict': model_restored.state_dict(),
'optimizer': optimizer.state_dict()
}, os.path.join(model_dir, "model_latest.pth"))
writer.add_scalar('train/loss', epoch_loss, epoch)
writer.add_scalar('train/lr', scheduler.get_lr()[0], epoch)
writer.close()
total_finish_time = (time.time() - total_start_time) # seconds
print('Total training time: {:.1f} hours'.format((total_finish_time / 60 / 60)))