Use accelerate

rct_diffusion_pipeline.py

@@ -30,7 +30,7 @@ class RCTDiffusionPipeline(DiffusionPipeline):
                               up_block_types=('UpBlock2D', 'CrossAttnUpBlock2D', 'CrossAttnUpBlock2D'), cross_attention_dim=160,
                             block_out_channels=(64, 128, 256), norm_num_groups=32)
 
-        self.unet.to(device='cuda', dtype=torch.float16)
+        self.unet.to(dtype=torch.float16)
     
     def load_dictionaries_from_dataset(self):
         dataset = load_dataset('frutiemax/rct_dataset')

train_model.py

@@ -9,6 +9,7 @@ import torchvision.transforms as T
 import torch.nn.functional as F
 from diffusers.optimization import get_cosine_schedule_with_warmup
 from tqdm.auto import tqdm
+from accelerate import Accelerator
 
 def save_and_test(pipeline, epoch):
     outputs = pipeline([[('aleppo pine tree', 1.0)]], [[('dark green', 1.0)]])
@@ -19,7 +20,7 @@ def save_and_test(pipeline, epoch):
     model_file = f'rct_foliage_{epoch}.pth'
     pipeline.save_pretrained(model_file)
 
-def train_model(batch_size=4, epochs=100, save_model_interval=10, start_learning_rate=1e-3, lr_warmup_steps=500):
+def train_model(batch_size=4, epochs=100, scheduler_num_timesteps=20, save_model_interval=10, start_learning_rate=1e-3, lr_warmup_steps=500):
     dataset = load_dataset('frutiemax/rct_dataset')
     dataset = dataset['train']
 
@@ -50,12 +51,12 @@ def train_model(batch_size=4, epochs=100, save_model_interval=10, start_learning
     del views
     
     # convert those views in tensors
-    targets = torch.Tensor(size=(num_images, 4, 3, 256, 256))
+    targets = torch.Tensor(size=(num_images, 4, 3, 256, 256)).to(dtype=torch.float16)
     pillow_to_tensor = T.ToTensor()
 
     for image_index in range(num_images):
         for view_index in range(4):
-            targets[image_index, view_index] = pillow_to_tensor(image_views[view_index][image_index])
+            targets[image_index, view_index] = pillow_to_tensor(image_views[view_index][image_index]).to(dtype=torch.float16)
     del image_views
     del entries
 
@@ -100,33 +101,50 @@ def train_model(batch_size=4, epochs=100, save_model_interval=10, start_learning
 
     # lets train for 100 epoch for each sprite in the dataset with a random noise level
     progress_bar = tqdm(total=epochs)
+    accelerator = Accelerator(
+        mixed_precision='fp16',
+        gradient_accumulation_steps=1,
+        log_with="tensorboard",
+        project_dir='logs',
+    )
+
+    unet, scheduler, optimizer, lr_scheduler = accelerator.prepare(model.unet, model.scheduler, \
+      optimizer, lr_scheduler)
+
+    del model
+    scheduler.set_timesteps(scheduler_num_timesteps)
     
     for epoch in range(epochs):
         # create a noisy version of each sprite
         for batch_index in range(0, num_images, batch_size):
             progress_bar.set_description(f'epoch={epoch}, batch_index={batch_index}')
             batch_end = np.minimum(num_images, batch_index + batch_size)
-            clean_images = targets[batch_index:batch_end].to(device='cuda', dtype=torch.float16)
+            clean_images = targets[batch_index:batch_end]
             clean_images = torch.reshape(clean_images, ((batch_end - batch_index), 12, 256, 256))
 
-            noise = torch.randn(clean_images.shape).to(device='cuda', dtype=torch.float16)
-            timesteps = torch.randint(0, model.scheduler.config.num_train_timesteps, (batch_end - batch_index, ))
-            timesteps = timesteps.to(dtype=torch.int, device='cuda')
-            noisy_images = model.scheduler.add_noise(clean_images, noise, timesteps).to(device='cuda', dtype=torch.float16)
-            noise_pred = model.unet(noisy_images, timesteps, class_labels[batch_index:batch_end].to(device='cuda',dtype=torch.float16), return_dict=False)[0]
+            noise = torch.randn(clean_images.shape, dtype=torch.float16)
+            timesteps = torch.randint(0, scheduler.config.num_train_timesteps, (batch_end - batch_index, ))
+            #timesteps = timesteps.to(dtype=torch.int, device='cuda')
+            noisy_images = scheduler.add_noise(clean_images, noise, timesteps).to(device='cuda', dtype=torch.float16)
+
+            with accelerator.accumulate(unet):
+                noise_pred = unet(noisy_images, timesteps.to(device='cuda', dtype=torch.float16), class_labels[batch_index:batch_end].to(device='cuda',dtype=torch.float16), return_dict=False)[0]
 
-            noise_pred = noise_pred.to(device='cuda', dtype=torch.float16)
-            loss = F.mse_loss(noise_pred, noise).to(device='cuda', dtype=torch.float16)
-            loss.backward()
+                #noise_pred = noise_pred.to(device='cuda', dtype=torch.float16)
+                loss = F.mse_loss(noise_pred, noise.to('cuda', dtype=torch.float16))
+                accelerator.backward(loss)
+                accelerator.clip_grad_norm_(unet.parameters(), 1.0)
 
-            optimizer.step()
-            lr_scheduler.step()
-            optimizer.zero_grad()
+                optimizer.step()
+                lr_scheduler.step()
+                optimizer.zero_grad()
         
         if (epoch + 1) % save_model_interval == 0:
+            model.unet = accelerator.unwrap_model(unet)
+            model.scheduler = scheduler
             save_and_test(model, epoch)
         progress_bar.update(1)
 
 
 if __name__ == '__main__':
-    train_model(8)
+    train_model(4)