updated tta to increase batch size in tta

.gitignore

@@ -4,6 +4,7 @@
 *.ipynb
 train_zero.py
 train_test.py
+test_time_adaptation_test.py
 /__pycache__
 /data
 /utils/__pycache__

config/adapt_config.py

@@ -55,4 +55,7 @@
 # Resource optimization flag. 0: intermediate files are saved, 1: intermediate files are deleted
 adapt_parser.add_argument('--resource', type=int, default=0, help=argparse.SUPPRESS)
 
+# batch size multiplier
+adapt_parser.add_argument('--batch_mul', type=int, default=4, help=argparse.SUPPRESS)
+
 args = adapt_parser.parse_args()

test_time_adaptation.py

@@ -4,7 +4,7 @@
 Test time adpatation module
 
 Editor: Marshall Xu
-Last Edited: 10/08/2023
+Last Edited: 08/10/2023
 """
 
 import os
@@ -125,10 +125,16 @@
         aug_item = aug_utils(patch_size, aug_mode)
 
         print("Finetuning procedure starts!")
+        print(f"\nIn this test, the batch size is {6*args.batch_mul}\n")
         # training loop
         for epoch in tqdm(range(epoch_num)):
             image, label = next(iter(data_loader))
             image_batch, label_batch = aug_item(image, label)
+            for j in range(1, args.batch_mul):
+                image, label = next(iter(data_loader))
+                image_batch_temp, label_batch_temp = aug_item(image, label)
+                image_batch = torch.cat((image_batch, image_batch_temp), 0)
+                label_batch = torch.cat((label_batch, label_batch_temp), 0)
             image_batch, label_batch = image_batch.to(device), label_batch.to(device)
 
             optimizer.zero_grad()
@@ -144,9 +150,9 @@
             # Learning rate shceduler
             scheduler.step(loss)
 
-            # TODO: debug message, delete this
             current_lr = optimizer.param_groups[0]['lr']
-            tqdm.write(f'Epoch: [{epoch+1}/{epoch_num}], Loss: {loss.item(): .4f}, Current learning rate: {current_lr: .8f}')
+
+        tqdm.write(f'Epoch: [{epoch+1}/{epoch_num}], Loss: {loss.item(): .4f}, Current learning rate: {current_lr : .8f}')
 
         file_name = processed_data_list[i].split('.')[0]
         out_mo_name = os.path.join(out_mo_path, file_name)

utils/unet_utils.py

@@ -341,6 +341,7 @@ class RandomCrop3D():
     """
     def __init__(self, img_sz, exp_sz):
         h, w, d = img_sz
+        # test 0925, constraint the higher bound of the crop size to be 128
         crop_h = torch.randint(32, h, (1,)).item()
         crop_w = torch.randint(32, w, (1,)).item()
         crop_d = torch.randint(32, d, (1,)).item()