Upload lowlight_train.py

lowlight_train.py

@@ -0,0 +1,124 @@
+import torch
+import torch.nn as nn
+import torchvision
+import torch.backends.cudnn as cudnn
+import torch.optim
+import os
+import sys
+import argparse
+import time
+import dataloader
+import model
+import Myloss
+import numpy as np
+from torchvision import transforms
+
+
+def weights_init(m):
+    classname = m.__class__.__name__
+    if classname.find('Conv') != -1:
+        m.weight.data.normal_(0.0, 0.02)
+    elif classname.find('BatchNorm') != -1:
+        m.weight.data.normal_(1.0, 0.02)
+        m.bias.data.fill_(0)
+
+
+
+
+
+def train(config):
+
+	os.environ['CUDA_VISIBLE_DEVICES']='0'
+
+	DCE_net = model.enhance_net_nopool().cuda()
+
+	DCE_net.apply(weights_init)
+	if config.load_pretrain == True:
+	    DCE_net.load_state_dict(torch.load(config.pretrain_dir))
+	train_dataset = dataloader.lowlight_loader(config.lowlight_images_path)		
+	
+	train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=config.train_batch_size, shuffle=True, num_workers=config.num_workers, pin_memory=True)
+
+
+
+	L_color = Myloss.L_color()
+	L_spa = Myloss.L_spa()
+
+	L_exp = Myloss.L_exp(16,0.6)
+	L_TV = Myloss.L_TV()
+
+
+	optimizer = torch.optim.Adam(DCE_net.parameters(), lr=config.lr, weight_decay=config.weight_decay)
+	
+	DCE_net.train()
+
+	for epoch in range(config.num_epochs):
+		for iteration, img_lowlight in enumerate(train_loader):
+
+			img_lowlight = img_lowlight.cuda()
+
+			enhanced_image_1,enhanced_image,A  = DCE_net(img_lowlight)
+
+			Loss_TV = 200*L_TV(A)
+			
+			loss_spa = torch.mean(L_spa(enhanced_image, img_lowlight))
+
+			loss_col = 5*torch.mean(L_color(enhanced_image))
+
+			loss_exp = 10*torch.mean(L_exp(enhanced_image))
+			
+			
+			# best_loss
+			loss =  Loss_TV + loss_spa + loss_col + loss_exp
+			#
+
+			
+			optimizer.zero_grad()
+			loss.backward()
+			torch.nn.utils.clip_grad_norm(DCE_net.parameters(),config.grad_clip_norm)
+			optimizer.step()
+
+			if ((iteration+1) % config.display_iter) == 0:
+				print("Loss at iteration", iteration+1, ":", loss.item())
+			if ((iteration+1) % config.snapshot_iter) == 0:
+				
+				torch.save(DCE_net.state_dict(), config.snapshots_folder + "Epoch" + str(epoch) + '.pth') 		
+
+
+
+
+if __name__ == "__main__":
+
+	parser = argparse.ArgumentParser()
+
+	# Input Parameters
+	parser.add_argument('--lowlight_images_path', type=str, default="data/train_data/")
+	parser.add_argument('--lr', type=float, default=0.0001)
+	parser.add_argument('--weight_decay', type=float, default=0.0001)
+	parser.add_argument('--grad_clip_norm', type=float, default=0.1)
+	parser.add_argument('--num_epochs', type=int, default=200)
+	parser.add_argument('--train_batch_size', type=int, default=8)
+	parser.add_argument('--val_batch_size', type=int, default=4)
+	parser.add_argument('--num_workers', type=int, default=4)
+	parser.add_argument('--display_iter', type=int, default=10)
+	parser.add_argument('--snapshot_iter', type=int, default=10)
+	parser.add_argument('--snapshots_folder', type=str, default="snapshots/")
+	parser.add_argument('--load_pretrain', type=bool, default= False)
+	parser.add_argument('--pretrain_dir', type=str, default= "snapshots/Epoch99.pth")
+
+	config = parser.parse_args()
+
+	if not os.path.exists(config.snapshots_folder):
+		os.mkdir(config.snapshots_folder)
+
+
+	train(config)
+
+
+
+
+
+
+
+
+