Update train.py

train.py

@@ -5,35 +5,30 @@ import random
 import datetime
 import argparse
 import numpy as np
-from tqdm import tqdm
-from piq import ssim,psnr
 from itertools import cycle
 
 import torch
 import torch.nn as nn
 from torch.utils import data
-import torch.distributed as dist 
-from torch.utils.data.distributed import DistributedSampler 
-from torch.nn.parallel import DistributedDataParallel as DDP 
 
+# Removed DDP and DistributedSampler imports
 
 from utils import dict2string,mkdir,get_lr,torch2cvimg,second2hours
-from loaders import docres_loader
-from models import restormer_arch
+# Assumed 'loaders' and 'models' modules are available
+from loaders import docres_loader 
+from models import restormer_arch 
 
+# --- Optional: Import for TensorBoard (uncomment if you have it installed) ---
+# from torch.utils.tensorboard import SummaryWriter
 
 def seed_torch(seed=1029):
 	random.seed(seed)
 	os.environ['PYTHONHASHSEED'] = str(seed) 
 	np.random.seed(seed)
 	torch.manual_seed(seed)
-	torch.cuda.manual_seed(seed)
-	torch.cuda.manual_seed_all(seed) 
+	# Removed CUDA-specific seeding
 	torch.backends.cudnn.benchmark = False
 	torch.backends.cudnn.deterministic = True
-    #torch.use_deterministic_algorithms(True)
-# seed_torch()
-
 
 def getBasecoord(h,w):
     base_coord0 = np.tile(np.arange(h).reshape(h,1),(1,w)).astype(np.float32)
@@ -42,13 +37,11 @@ def getBasecoord(h,w):
     return base_coord
 
 def train(args):
-
-    ## DDP init
-    dist.init_process_group(backend='nccl',init_method='env://',timeout=datetime.timedelta(seconds=36000))
-    torch.cuda.set_device(args.local_rank)
-    device = torch.device('cuda',args.local_rank)
-    torch.cuda.manual_seed_all(42)
-
+    # --- CPU/Single-Process Setup ---
+    # Set device to CPU
+    device = torch.device('cpu')
+    print(f"Training on device: {device}")
+    
     ### Log file:
     mkdir(args.logdir)
     mkdir(os.path.join(args.logdir,args.experiment_name))
@@ -58,10 +51,17 @@ def train(args):
     log_file.close()
 
     ### Setup tensorboard for visualization
-    if args.tboard:
-        writer = SummaryWriter(os.path.join(args.logdir,args.experiment_name,'runs'),args.experiment_name)
+    # Note: TensorBoard setup is commented out for robust CPU execution.
+    # if args.tboard:
+    #     try:
+    #         writer = SummaryWriter(os.path.join(args.logdir,args.experiment_name,'runs'),args.experiment_name)
+    #     except NameError:
+    #         print("Warning: TensorBoard not imported. Skipping logging to SummaryWriter.")
+    #         args.tboard = False
+
 
     ### Setup Dataloader
+    # NOTE: You MUST update these paths to match your system setup.
     datasets_setting = [
         {'task':'deblurring','ratio':1,'im_path':'/home/jiaxin/Training_Data/DocRes_data/train/deblurring/','json_paths':['/home/jiaxin/Training_Data/DocRes_data/train/deblurring/tdd/train.json']},
         {'task':'dewarping','ratio':1,'im_path':'/home/jiaxin/Training_Data/DocRes_data/train/dewarping/','json_paths':['/home/jiaxin/Training_Data/DocRes_data/train/dewarping/doc3d/train_1_19.json']},
@@ -69,33 +69,30 @@ def train(args):
         {'task':'deshadowing','ratio':1,'im_path':'/home/jiaxin/Training_Data/DocRes_data/train/deshadowing/','json_paths':['/home/jiaxin/Training_Data/DocRes_data/train/deshadowing/train.json']},
         {'task':'appearance','ratio':1,'im_path':'/home/jiaxin/Training_Data/DocRes_data/train/appearance/','json_paths':['/home/jiaxin/Training_Data/DocRes_data/train/appearance/trainv2.json']}
         ]
-
-
     ratios = [dataset_setting['ratio'] for dataset_setting in datasets_setting]
     datasets = [docres_loader.DocResTrainDataset(dataset=dataset_setting,img_size=args.im_size) for dataset_setting in datasets_setting]
-    trainloaders = [{'task':datasets_setting[i],'loader':data.DataLoader(dataset=datasets[i], sampler=DistributedSampler(datasets[i]), batch_size=args.batch_size, num_workers=2, pin_memory=True,drop_last=True),'iter_loader':iter(data.DataLoader(dataset=datasets[i], sampler=DistributedSampler(datasets[i]), batch_size=args.batch_size, num_workers=2, pin_memory=True,drop_last=True))} for i in range(len(datasets))]
-
-
-    ### test loader
-    # for i in tqdm(range(args.total_iter)):
-    #     loader_index = random.choices(list(range(len(trainloaders))),ratios)[0]
-    #     in_im,gt_im = next(trainloaders[loader_index]['iter_loader'])
-
-
+    
+    # Standard DataLoader is used instead of DistributedSampler
+    trainloaders = [{'task':datasets_setting[i],
+                     'loader':data.DataLoader(dataset=datasets[i], batch_size=args.batch_size, num_workers=0, pin_memory=False, drop_last=True),
+                     'iter_loader':iter(data.DataLoader(dataset=datasets[i], batch_size=args.batch_size, num_workers=0, pin_memory=False, drop_last=True))} 
+                    for i in range(len(datasets))]
+    
     ### Setup Model
-    model = restormer_arch.Restormer( 
-        inp_channels=6, 
-        out_channels=3, 
-        dim = 48,
-        num_blocks = [2,3,3,4],     
-        num_refinement_blocks = 4,
+    model = restormer_arch.Restormer(
+         inp_channels=6, 
+         out_channels=3, 
+         dim = 48,
+        num_blocks = [2,3,3,4], 
+            num_refinement_blocks = 4,
         heads = [1,2,4,8],
         ffn_expansion_factor = 2.66,
         bias = False,
-        LayerNorm_type = 'WithBias',   
-        dual_pixel_task = True       
-    )
-    model=DDP(model.cuda(),device_ids=[args.local_rank],output_device=args.local_rank)
+        LayerNorm_type = 'WithBias',
+        dual_pixel_task = True           
+           )
+    # Move model to CPU
+    model.to(device)
 
     ### Optimizer
     optimizer= torch.optim.AdamW(model.parameters(),lr=args.l_rate,weight_decay=5e-4)
@@ -105,8 +102,10 @@ def train(args):
 
     ### load checkpoint
     iter_start=0
-    if args.resume is not None:                                         
+    if args.resume is not None:                                                 
         print("Loading model and optimizer from checkpoint '{}'".format(args.resume))
+        # Ensure checkpoint is loaded to CPU
+        checkpoint = torch.load(args.resume, map_location=device)
         x = checkpoint['model_state']
         model.load_state_dict(x,strict=False)
         iter_start=checkpoint['iter']
@@ -114,7 +113,7 @@ def train(args):
 
     ###-----------------------------------------Training-----------------------------------------
     ##initialize
-    scaler = torch.cuda.amp.GradScaler()
+    # Removed GradScaler for AMP
     loss_dict = {}
     total_step = 0
     l2 = nn.MSELoss()
@@ -128,74 +127,80 @@ def train(args):
     ## total_steps
     for iters in range(iter_start,args.total_iter):
         start_time = time.time()
+        
         loader_index = random.choices(list(range(len(trainloaders))),ratios)[0]
-
         try:
             in_im,gt_im = next(trainloaders[loader_index]['iter_loader'])
         except StopIteration:
             trainloaders[loader_index]['iter_loader']=iter(trainloaders[loader_index]['loader'])
             in_im,gt_im = next(trainloaders[loader_index]['iter_loader'])
-        in_im = in_im.float().cuda()
-        gt_im = gt_im.float().cuda()
-
+        
+        # Move data to CPU
+        in_im = in_im.float().to(device)
+        gt_im = gt_im.float().to(device)
+        
         binarization_loss,appearance_loss,dewarping_loss,deblurring_loss,deshadowing_loss = 0,0,0,0,0
-        with torch.cuda.amp.autocast():
-            pred_im = model(in_im,trainloaders[loader_index]['task']['task'])
-            if trainloaders[loader_index]['task']['task'] == 'binarization':
-                gt_im = gt_im.long()
-                binarization_loss = ce(pred_im[:,:2,:,:], gt_im[:,0,:,:])
-                loss = binarization_loss
-            elif trainloaders[loader_index]['task']['task'] == 'dewarping':
-                dewarping_loss = l1(pred_im[:,:2,:,:], gt_im[:,:2,:,:])
-                loss = dewarping_loss
-            elif trainloaders[loader_index]['task']['task'] == 'appearance':
-                appearance_loss = l1(pred_im, gt_im)
-                loss = appearance_loss
-            elif trainloaders[loader_index]['task']['task'] == 'deblurring':
-                deblurring_loss = l1(pred_im, gt_im)
-                loss = deblurring_loss
-            elif trainloaders[loader_index]['task']['task'] == 'deshadowing':
-                deshadowing_loss = l1(pred_im, gt_im)
-                loss = deshadowing_loss
 
+        # Removed torch.cuda.amp.autocast() block
+        pred_im = model(in_im,trainloaders[loader_index]['task']['task'])
+        
+        if trainloaders[loader_index]['task']['task'] == 'binarization':
+            gt_im = gt_im.long()
+            binarization_loss = ce(pred_im[:,:2,:,:], gt_im[:,0,:,:])
+            loss = binarization_loss
+        elif trainloaders[loader_index]['task']['task'] == 'dewarping':
+            dewarping_loss = l1(pred_im[:,:2,:,:], gt_im[:,:2,:,:])
+            loss = dewarping_loss
+        elif trainloaders[loader_index]['task']['task'] == 'appearance':
+            appearance_loss = l1(pred_im, gt_im)
+            loss = appearance_loss
+        elif trainloaders[loader_index]['task']['task'] == 'deblurring':
+            deblurring_loss = l1(pred_im, gt_im)
+            loss = deblurring_loss
+        elif trainloaders[loader_index]['task']['task'] == 'deshadowing':
+            deshadowing_loss = l1(pred_im, gt_im)
+            loss = deshadowing_loss
+        
         optimizer.zero_grad()
-        scaler.scale(loss).backward()
-        scaler.step(optimizer)
-        scaler.update() 
-    
+        # Standard backward pass (removed scaler)
+        loss.backward()
+        optimizer.step()
+        
         loss_dict['dew_loss']=dewarping_loss.item() if isinstance(dewarping_loss,torch.Tensor) else 0
         loss_dict['app_loss']=appearance_loss.item() if isinstance(appearance_loss,torch.Tensor) else 0
         loss_dict['des_loss']=deshadowing_loss.item() if isinstance(deshadowing_loss,torch.Tensor) else 0
         loss_dict['deb_loss']=deblurring_loss.item() if isinstance(deblurring_loss,torch.Tensor) else 0
         loss_dict['bin_loss']=binarization_loss.item() if isinstance(binarization_loss,torch.Tensor) else 0
+
         end_time = time.time()
         duration = end_time-start_time
+        
         ## log
         if (iters+1) % 10 == 0:
             ## print
             print('iters [{}/{}] -- '.format(iters+1,args.total_iter)+dict2string(loss_dict)+' --lr {:6f}'.format(get_lr(optimizer))+' -- time {}'.format(second2hours(duration*(args.total_iter-iters))))
+            
             ## tbord
-            if args.tboard:
-                for key,value in loss_dict.items():
-                    writer.add_scalar('Train '+key+'/Iterations', value, total_step)
+            # if args.tboard:
+            #     for key,value in loss_dict.items():
+            #         writer.add_scalar('Train '+key+'/Iterations', value, total_step)
+            
             ## logfile
             with open(log_file_path,'a') as f:
                 f.write('iters [{}/{}] -- '.format(iters+1,args.total_iter)+dict2string(loss_dict)+' --lr {:6f}'.format(get_lr(optimizer))+' -- time {}'.format(second2hours(duration*(args.total_iter-iters)))+'\n')
 
-
         if (iters+1) % 5000 == 0:
             state = {'iters': iters+1,
                      'model_state': model.state_dict(),
                      'optimizer_state' : optimizer.state_dict(),}
-            if not os.path.exists(os.path.join(args.logdir,args.experiment_name)):
+            if not os.path.exists(os.path.join(args.logdir,args.experiment_name)): 
                  os.system('mkdir ' + os.path.join(args.logdir,args.experiment_name))
-            if torch.distributed.get_rank()==0:
-                torch.save(state, os.path.join(args.logdir,args.experiment_name,"{}.pkl".format(iters+1)))
-
+            
+            # Save checkpoint without DDP rank check
+            torch.save(state, os.path.join(args.logdir,args.experiment_name,"{}.pkl".format(iters+1)))
+            
         sched.step()
 
-
-
 if __name__ == '__main__':
     parser = argparse.ArgumentParser(description='Hyperparams')
     parser.add_argument('--im_size', nargs='?', type=int, default=256, 
@@ -206,16 +211,19 @@ if __name__ == '__main__':
                         help='Batch Size')
     parser.add_argument('--l_rate', nargs='?', type=float, default=2e-4, 
                         help='Learning Rate')
-    parser.add_argument('--resume', nargs='?', type=str, default=None,    
-                        help='Path to previous saved model to restart from')
-    parser.add_argument('--logdir', nargs='?', type=str, default='./checkpoints/',    
-                        help='Path to store the loss logs')
+    parser.add_argument('--resume', nargs='?', type=str, default=None,
+                            help='Path to previous saved model to restart from')
+    parser.add_argument('--logdir', nargs='?', type=str, default='./checkpoints/',
+                            help='Path to store the loss logs')
     parser.add_argument('--tboard', dest='tboard', action='store_true', 
                         help='Enable visualization(s) on tensorboard | False by default')
-    parser.add_argument('--local_rank',type=int,default=0,metavar='N')    
+    # Removed local_rank argument as it's not needed for single-process CPU
     parser.add_argument('--experiment_name', nargs='?', type=str,default='experiment_name',
                         help='the name of this experiment')
     parser.set_defaults(tboard=False)
     args = parser.parse_args()
-
-    train(args)
+    
+    # Note: Using a low batch size (e.g., 2) is recommended for initial CPU testing.
+    # args.batch_size = 2 # Uncomment for quick testing
+    
+    train(args)