Update glide_text2im/train_util.py

glide_text2im/train_util.py

@@ -9,7 +9,7 @@ import torch.distributed as dist
 from torch.nn.parallel.distributed import DistributedDataParallel as DDP
 from torch.optim import AdamW
 from .glide_util import sample
-from . import dist_util, logger
+from . import  logger
 from .fp16_util import (
     make_master_params,
     master_params_to_model_params,
@@ -84,7 +84,7 @@ class TrainLoop:
         self.global_batch = self.batch_size * dist.get_world_size()
 
         if use_vgg:
-            self.vgg = VGG(conv_index='22').to(dist_util.dev())
+            self.vgg = VGG(conv_index='22').cuda()
             print('use perc')
         else:
             self.vgg = None
@@ -131,8 +131,8 @@ class TrainLoop:
             self.use_ddp = True
             self.ddp_model = DDP(
                 self.model,
-                device_ids=[dist_util.dev()],
-                output_device=dist_util.dev(),
+                device_ids=[torch.device('cuda')],
+                output_device=torch.device('cuda'),
                 broadcast_buffers=False,
                 bucket_cap_mb=128,
                 find_unused_parameters=False,
@@ -155,7 +155,7 @@ class TrainLoop:
                 logger.log(f"loading model from checkpoint: {resume_checkpoint}...")
                 self.model.load_state_dict(th.load(resume_checkpoint, map_location="cpu"),strict=False)
 
-        dist_util.sync_params(self.model.parameters())
+        #dist_util.sync_params(self.model.parameters())
 
     def _load_ema_parameters(self, rate):
         ema_params = copy.deepcopy(self.master_params)
@@ -165,7 +165,7 @@ class TrainLoop:
         if ema_checkpoint:
             if dist.get_rank() == 0:
                 logger.log(f"loading EMA from checkpoint: {ema_checkpoint}...")
-                state_dict = th.load(ema_checkpoint, map_location=dist_util.dev())
+                state_dict = th.load(ema_checkpoint, map_location=torch.device('cuda'))
                 ema_params = self._state_dict_to_master_params(state_dict)
 
         #dist_util.sync_params(ema_params)
@@ -227,10 +227,10 @@ class TrainLoop:
     def forward_backward(self, batch, model_kwargs):
         zero_grad(self.model_params)
         for i in range(0, batch.shape[0], self.microbatch):
-            micro = batch[i : i + self.microbatch].to(dist_util.dev())
-            micro_cond={n:model_kwargs[n][i:i+self.microbatch].to(dist_util.dev()) for n in model_kwargs if n  in ['ref', 'low_res']}
+            micro = batch[i : i + self.microbatch].cuda()
+            micro_cond={n:model_kwargs[n][i:i+self.microbatch].cuda() for n in model_kwargs if n  in ['ref', 'low_res']}
             last_batch = (i + self.microbatch) >= batch.shape[0]
-            t, weights = self.schedule_sampler.sample(micro.shape[0], dist_util.dev())
+            t, weights = self.schedule_sampler.sample(micro.shape[0], torch.device('cuda'))
              
             if self.step <100:
                 vgg_loss = None
@@ -295,7 +295,7 @@ class TrainLoop:
                     prompt=model_kwargs,
                     batch_size=self.glide_options['batch_size']//2,
                     guidance_scale=guidance_scale,
-                    device=dist_util.dev(),
+                    device=torch.device('cuda'),
                     prediction_respacing=self.glide_options['sample_respacing'],
                     upsample_enabled=self.glide_options['super_res'],
                     upsample_temp=0.997,