text: 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 net
  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):
  
  	dehaze_net = net.dehaze_net().cuda()
  	dehaze_net.apply(weights_init)
  
  	train_dataset = dataloader.dehazing_loader(config.orig_images_path,
  											 config.hazy_images_path)		
  	val_dataset = dataloader.dehazing_loader(config.orig_images_path,
  											 config.hazy_images_path, mode="val")		
  	train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=config.train_batch_size, shuffle=True, num_workers=config.num_workers, pin_memory=True)
  	val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=config.val_batch_size, shuffle=True, num_workers=config.num_workers, pin_memory=True)
  
  	criterion = nn.MSELoss().cuda()
  	optimizer = torch.optim.Adam(dehaze_net.parameters(), lr=config.lr, weight_decay=config.weight_decay)
  	
  	dehaze_net.train()
  
  	for epoch in range(config.num_epochs):
  		for iteration, (img_orig, img_haze) in enumerate(train_loader):
  
  			img_orig = img_orig.cuda()
  			img_haze = img_haze.cuda()
  
  			clean_image = dehaze_net(img_haze)
  
  			loss = criterion(clean_image, img_orig)
  
  			optimizer.zero_grad()
  			loss.backward()
  			torch.nn.utils.clip_grad_norm(dehaze_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(dehaze_net.state_dict(), config.snapshots_folder + "Epoch" + str(epoch) + '.pth') 		
  
  		# Validation Stage
  		for iter_val, (img_orig, img_haze) in enumerate(val_loader):
  
  			img_orig = img_orig.cuda()
  			img_haze = img_haze.cuda()
  
  			clean_image = dehaze_net(img_haze)
  
  			torchvision.utils.save_image(torch.cat((img_haze, clean_image, img_orig),0), config.sample_output_folder+str(iter_val+1)+".jpg")
  
  		torch.save(dehaze_net.state_dict(), config.snapshots_folder + "dehazer.pth") 
  
  			
  
  
  
  
  
  
  
  
  if __name__ == "__main__":
  
  	parser = argparse.ArgumentParser()
  
  	# Input Parameters
  	parser.add_argument('--orig_images_path', type=str, default="data/images/")
  	parser.add_argument('--hazy_images_path', type=str, default="data/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=10)
  	parser.add_argument('--train_batch_size', type=int, default=8)
  	parser.add_argument('--val_batch_size', type=int, default=8)
  	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=200)
  	parser.add_argument('--snapshots_folder', type=str, default="snapshots/")
  	parser.add_argument('--sample_output_folder', type=str, default="samples/")
  
  	config = parser.parse_args()
  
  	if not os.path.exists(config.snapshots_folder):
  		os.mkdir(config.snapshots_folder)
  	if not os.path.exists(config.sample_output_folder):
  		os.mkdir(config.sample_output_folder)
  
  	train(config)
  
  
  
  
  
  
  
  
  	
  