divide images

HowieMa · Feb 10, 2019 · b0d982f · b0d982f
1 parent 0b00cb9
commit b0d982f
Show file tree

Hide file tree

Showing 6 changed files with 71 additions and 70 deletions.
diff --git a/.idea/workspace.xml b/.idea/workspace.xml
diff --git a/model/stacked_unet3d.py b/model/stacked_unet3d.py
@@ -34,8 +34,8 @@ def forward(self, x):
 
 # test case
 if __name__ == "__main__":
-    net = StackedUnet3D(in_ch=1, out_ch=2, degree=16)
-    print"total parameter:" + str(netSize(net))  # 25906856
+    net = StackedUnet3D(in_ch=1, out_ch=2, degree=8)
+    print"total parameter:" + str(netSize(net))  # 6483800
 
     x = torch.randn(4, 4, 16, 192, 192)  # batch size = 2
     print ('input data')

diff --git a/model/unet3d.py b/model/unet3d.py
@@ -66,8 +66,8 @@ def forward(self, x):
 
 # test case
 if __name__ == "__main__":
-    net = UNet3D(4, 2, degree=32)
-    print"total parameter:" + str(netSize(net))     # 25893986
+    net = UNet3D(4, 2, degree=16)
+    print"total parameter:" + str(netSize(net))     # 6477362 25MB
 
     x = torch.randn(4, 4, 16, 192, 192)  # batch size = 2
     print ('input data')

diff --git a/model/unet3d.pyc b/model/unet3d.pyc
diff --git a/train.py b/train.py
@@ -32,7 +32,7 @@
 
 # ******************** build model ********************
 if model == '3dunet':
-    net = UNet3D(in_ch=4, out_ch=2, degree=32)  # multi-modal =4, out binary classification one-hot
+    net = UNet3D(in_ch=4, out_ch=2, degree=16)  # multi-modal =4, out binary classification one-hot
 elif model == 'multi_unet':
     net = Multi_Unet(1, 2, 32)
 else:
@@ -97,7 +97,7 @@ def run():
                     # 5D float Tensor   Batch_Size * 2 * 16(volume_size) * height * weight
                     predicts = (predicts[:, 1, :, :, :] > 0.5).long()
                     # 4D Long  Tensor   Batch_Size * 16(volume_size) * height * weight
-                    d = dice(predicts, labels[:, 0, :, :, :].long())
+                    d = dice(predicts, lbl[:, 0, :, :, :].long())
                     train_dice.append(d)
 
                 # ****** save image of step 0 for each epoch ******
@@ -127,7 +127,7 @@ def run():
                     # 5D float Tensor   Batch_Size * 2 * 16(volume_size) * height * weight
                     predicts = (predicts[:, 1, :, :, :] > 0.5).long()
                     # 4D Long  Tensor   Batch_Size * 16(volume_size) * height * weight
-                    d = dice(predicts, labels[:, 0, :, :, :].long())
+                    d = dice(predicts, lbl[:, 0, :, :, :].long())
                     test_dice.append(d)
 
         # **************** save loss for one batch ****************

diff --git a/train_stacked.py b/train_stacked.py
@@ -107,7 +107,7 @@ def run():
                 labels = Variable(labels_vol[i].cuda() if cuda_available else labels_vol[i])
                 # 5D tensor   Batch_Size * 1        * 16 * 192 * 192
 
-                
+
 
                 predicts = net(images)
                 # # 6D tensor   Batch_Size * 4 * 2 * 16(volume_size) * height * weight