updated

Author: dhchenx

GANs/3dGAN/src/__pycache__/model.cpython-36.pyc

@@ -0,0 +1,50 @@
+
+'''
+main.py
+
+Welcome, this is the entrance to 3dgan
+'''
+import os
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+import argparse
+from trainer import trainer
+import torch
+
+from tester import tester
+import params
+
+def str2bool(v):
+    if v.lower() in ('yes', 'true', 't', 'y', '1'):
+        return True
+    elif v.lower() in ('no', 'false', 'f', 'n', '0'):
+        return False
+    else:
+        raise argparse.ArgumentTypeError('Boolean value expected.')
+
+def main():
+
+    # add arguments
+    parser = argparse.ArgumentParser()
+
+    # loggings parameters
+    parser.add_argument('--logs', type=str, default=None, help='logs by tensorboardX')
+    parser.add_argument('--local_test', type=str2bool, default=False, help='local test verbose')
+    parser.add_argument('--model_name', type=str, default="dcgan", help='model name for saving')
+    parser.add_argument('--test', type=str2bool, default=False, help='call tester.py')
+    parser.add_argument('--use_visdom', type=str2bool, default=False, help='visualization by visdom')
+    args = parser.parse_args()
+
+    # list params
+    params.print_params()
+
+    # run program
+    if args.test == False:
+        trainer(args)
+    else:
+        tester(args)
+
+
+if __name__ == '__main__':
+    main()
+
+    

GANs/3dGAN/src/__pycache__/params.cpython-36.pyc

@@ -0,0 +1,120 @@
+import torch
+import params
+
+'''
+
+model.py
+
+Define our GAN model
+
+The cube_len is 32x32x32, and the maximum number of feature map is 256, 
+so the results may be inconsistent with the paper
+
+'''
+
+class net_G(torch.nn.Module):
+    def __init__(self, args):
+        super(net_G, self).__init__()
+        self.args = args
+        self.cube_len = params.cube_len
+        self.bias = params.bias
+        self.z_dim = params.z_dim
+        self.f_dim = 32
+
+        padd = (0, 0, 0)
+        if self.cube_len == 32:
+            padd = (1,1,1)
+
+        self.layer1 = self.conv_layer(self.z_dim, self.f_dim*8, kernel_size=4, stride=2, padding=padd, bias=self.bias)
+        self.layer2 = self.conv_layer(self.f_dim*8, self.f_dim*4, kernel_size=4, stride=2, padding=(1, 1, 1), bias=self.bias)
+        self.layer3 = self.conv_layer(self.f_dim*4, self.f_dim*2, kernel_size=4, stride=2, padding=(1, 1, 1), bias=self.bias)
+        self.layer4 = self.conv_layer(self.f_dim*2, self.f_dim, kernel_size=4, stride=2, padding=(1, 1, 1), bias=self.bias)
+        
+        self.layer5 = torch.nn.Sequential(
+            torch.nn.ConvTranspose3d(self.f_dim, 1, kernel_size=4, stride=2, bias=self.bias, padding=(1, 1, 1)),
+            torch.nn.Sigmoid()
+            # torch.nn.Tanh()
+        )
+
+    def conv_layer(self, input_dim, output_dim, kernel_size=4, stride=2, padding=(1,1,1), bias=False):
+        layer = torch.nn.Sequential(
+            torch.nn.ConvTranspose3d(input_dim, output_dim, kernel_size=kernel_size, stride=stride, bias=bias, padding=padding),
+            torch.nn.BatchNorm3d(output_dim),
+            torch.nn.ReLU(True)
+            # torch.nn.LeakyReLU(self.leak_value, True)
+        )
+        return layer
+
+    def forward(self, x):
+        out = x.view(-1, self.z_dim, 1, 1, 1)
+        # print(out.size())  # torch.Size([32, 200, 1, 1, 1])
+        out = self.layer1(out)
+        # print(out.size())  # torch.Size([32, 256, 2, 2, 2])
+        out = self.layer2(out)
+        # print(out.size())  # torch.Size([32, 128, 4, 4, 4])
+        out = self.layer3(out)
+        # print(out.size())  # torch.Size([32, 64, 8, 8, 8])
+        out = self.layer4(out)
+        # print(out.size())  # torch.Size([32, 32, 16, 16, 16])
+        out = self.layer5(out)
+        # print(out.size())  # torch.Size([32, 1, 32, 32, 32])
+        out = torch.squeeze(out)
+        return out
+
+
+class net_D(torch.nn.Module):
+    def __init__(self, args):
+        super(net_D, self).__init__()
+        self.args = args
+        self.cube_len = params.cube_len
+        self.leak_value = params.leak_value
+        self.bias = params.bias
+
+        padd = (0,0,0)
+        if self.cube_len == 32:
+            padd = (1,1,1)
+
+        self.f_dim = 32
+
+        self.layer1 = self.conv_layer(1, self.f_dim, kernel_size=4, stride=2, padding=(1,1,1), bias=self.bias)
+        self.layer2 = self.conv_layer(self.f_dim, self.f_dim*2, kernel_size=4, stride=2, padding=(1,1,1), bias=self.bias)
+        self.layer3 = self.conv_layer(self.f_dim*2, self.f_dim*4, kernel_size=4, stride=2, padding=(1,1,1), bias=self.bias)
+        self.layer4 = self.conv_layer(self.f_dim*4, self.f_dim*8, kernel_size=4, stride=2, padding=(1,1,1), bias=self.bias)
+
+        self.layer5 = torch.nn.Sequential(
+            torch.nn.Conv3d(self.f_dim*8, 1, kernel_size=4, stride=2, bias=self.bias, padding=padd),
+            torch.nn.Sigmoid()
+        )
+
+        # self.layer5 = torch.nn.Sequential(
+        #     torch.nn.Linear(256*2*2*2, 1),
+        #     torch.nn.Sigmoid()
+        # )
+
+    def conv_layer(self, input_dim, output_dim, kernel_size=4, stride=2, padding=(1,1,1), bias=False):
+        layer = torch.nn.Sequential(
+            torch.nn.Conv3d(input_dim, output_dim, kernel_size=kernel_size, stride=stride, bias=bias, padding=padding),
+            torch.nn.BatchNorm3d(output_dim),
+            torch.nn.LeakyReLU(self.leak_value, inplace=True)
+        )
+        return layer
+
+    def forward(self, x):
+        # out = torch.unsqueeze(x, dim=1)
+        out = x.view(-1, 1, self.cube_len, self.cube_len, self.cube_len)
+        # print(out.size()) # torch.Size([32, 1, 32, 32, 32])
+        out = self.layer1(out)
+        # print(out.size())  # torch.Size([32, 32, 16, 16, 16])
+        out = self.layer2(out)
+        # print(out.size())  # torch.Size([32, 64, 8, 8, 8])
+        out = self.layer3(out)
+        # print(out.size())  # torch.Size([32, 128, 4, 4, 4])
+        out = self.layer4(out)
+        # print(out.size())  # torch.Size([32, 256, 2, 2, 2])
+        # out = out.view(-1, 256*2*2*2)
+        # print (out.size())
+        out = self.layer5(out)
+        # print(out.size())  # torch.Size([32, 1, 1, 1, 1])
+        out = torch.squeeze(out)
+        return out
+

GANs/3dGAN/src/__pycache__/tester.cpython-36.pyc

@@ -0,0 +1,59 @@
+
+'''
+params.py
+
+Managers of all hyper-parameters
+
+'''
+
+import torch
+
+epochs = 500
+batch_size = 32
+soft_label = False
+adv_weight = 0
+d_thresh = 0.8
+z_dim = 200
+z_dis = "norm"
+model_save_step = 1
+g_lr = 0.0025
+d_lr = 0.00001
+beta = (0.5, 0.999)
+cube_len = 32
+leak_value = 0.2
+bias = False
+
+device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+data_dir = '../volumetric_data/'
+model_dir = 'airplane/'    # change it to train on other data models
+output_dir = '../outputs'
+images_dir = '../images'
+
+def print_params():
+    l = 16
+    print (l*'*' + 'hyper-parameters' + l*'*')
+
+    print ('epochs =', epochs)
+    print ('batch_size =', batch_size)
+    print ('soft_labels =', soft_label)
+    print ('adv_weight =', adv_weight)
+    print ('d_thresh =', d_thresh)
+    print ('z_dim =', z_dim)
+    print ('z_dis =', z_dis)
+    print ('model_images_save_step =', model_save_step)
+    print ('data =', model_dir)
+    print ('device =', device)
+    print ('g_lr =', g_lr)
+    print ('d_lr =', d_lr)
+    print ('cube_len =', cube_len)
+    print ('leak_value =', leak_value)
+    print ('bias =', bias)
+
+    print (l*'*' + 'hyper-parameters' + l*'*')
+
+
+
+
+
+
+

GANs/3dGAN/src/__pycache__/trainer.cpython-36.pyc

@@ -0,0 +1,102 @@
+
+'''
+tester.py
+
+Test the trained 3dgan models
+'''
+
+import torch
+from torch import optim
+from torch import nn
+from collections import OrderedDict
+from utils import *
+import os
+from model import net_G, net_D
+# from lr_sh import  MultiStepLR
+
+# added
+import datetime
+from tensorboardX import SummaryWriter
+import matplotlib.pyplot as plt
+import numpy as np
+import params
+import visdom
+
+# def test_gen(args):
+#     test_z = []
+#     test_num = 1000
+#     for i in range(test_num):
+#         z = generateZ(args, 1)
+#         z = z.numpy()
+#         test_z.append(z)
+    
+#     test_z = np.array(test_z)
+#     print (test_z.shape)
+    # np.save("test_z", test_z)
+
+def tester(args):
+    print ('Evaluation Mode...')
+
+    # image_saved_path = '../images'
+    image_saved_path = params.images_dir
+    if not os.path.exists(image_saved_path):
+        os.makedirs(image_saved_path)
+
+    if args.use_visdom == True:
+        vis = visdom.Visdom()
+
+    save_file_path = params.output_dir + '/' + args.model_name
+    pretrained_file_path_G = save_file_path+'/'+'G.pth'
+    pretrained_file_path_D = save_file_path+'/'+'D.pth'
+    
+    print (pretrained_file_path_G)
+
+    D = net_D(args)
+    G = net_G(args)
+
+    if not torch.cuda.is_available():
+        G.load_state_dict(torch.load(pretrained_file_path_G, map_location={'cuda:0': 'cpu'}))
+        D.load_state_dict(torch.load(pretrained_file_path_D, map_location={'cuda:0': 'cpu'}))
+    else:
+        G.load_state_dict(torch.load(pretrained_file_path_G))
+        D.load_state_dict(torch.load(pretrained_file_path_D, map_location={'cuda:0': 'cpu'}))
+    
+    print ('visualizing model')
+    
+    # test generator
+    # test_gen(args)
+    G.to(params.device)
+    D.to(params.device)
+    G.eval()
+    D.eval()
+
+    # test_z = np.load("test_z.npy")
+    # print (test_z.shape)
+    # N = test_z.shape[0]
+
+    N = 8
+
+    for i in range(N):
+        # z = test_z[i,:]
+        # z = torch.FloatTensor(z)
+        
+        z = generateZ(args, 1)
+        
+        # print (z.size())
+        fake = G(z)
+        samples = fake.unsqueeze(dim=0).detach().cpu().numpy()
+        # print (samples.shape)
+        # print (fake)
+        y_prob = D(fake)
+        y_real = torch.ones_like(y_prob)
+        # criterion = nn.BCELoss()
+        # print (y_prob.item(), criterion(y_prob, y_real).item())
+
+        ### visualization
+        if args.use_visdom == False:
+            SavePloat_Voxels(samples, image_saved_path, 'tester_norm_'+str(i))
+        else:
+            plotVoxelVisdom(samples[0,:], vis, "tester_"+str(i))
+        
+
+