forked from EdwardSmith1884/3D-IWGAN
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
0 parents
commit 0575f30
Showing
95 changed files
with
182,533 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,123 @@ | ||
import tensorflow as tf | ||
import keras.backend as K | ||
import os | ||
import sys | ||
import tensorlayer as tl | ||
import numpy as np | ||
from tensorlayer.layers import * | ||
import random | ||
from glob import glob | ||
import argparse | ||
import scripts | ||
from scripts.GANutils import * | ||
from scripts.models import * | ||
|
||
|
||
|
||
|
||
parser = argparse.ArgumentParser(description='3D-GAN implementation for 32*32*32 voxel output') | ||
parser.add_argument('-n','--name', default='Test', help='The name of the current experiment, this will be used to create folders and save models.') | ||
parser.add_argument('-d','--data', default='data/train/chair', help ='The location fo the object voxel models.' ) | ||
parser.add_argument('-e','--epochs', default=1500, help ='The number of epochs to run for.', type=int) | ||
parser.add_argument('-b','--batchsize', default=256, help ='The batch size.', type=int) | ||
parser.add_argument('-sample', default= 5, help='How often generated obejcts are sampled and saved.', type= int) | ||
parser.add_argument('-save', default= 5, help='How often the network models are saved.', type= int) | ||
parser.add_argument('-l', '--load', default= False, help='Indicates if a previously loaded model should be loaded.', action = 'store_true') | ||
parser.add_argument('-le', '--load_epoch', default= '', help='The epoch to number to be loaded from.', type=str) | ||
parser.add_argument('-mbd', '--mini_batch_discimination', default= False, help= 'Indicates if Mini Bathc Discrimination should be used', action = 'store_true') | ||
parser.add_argument('-glr','--genorator_learning_rate', default=0.0025, help ='The genorator learning rate.', type=int) | ||
parser.add_argument('-dlr','--discriminator_learning_rate', default=0.00005, help ='The discriminator learning rate.', type=int) | ||
|
||
args = parser.parse_args() | ||
|
||
checkpoint_dir = "checkpoint/" + args.name +'/' | ||
save_dir = "savepoint/" + args.name +'/' | ||
output_size = 32 | ||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
######### make directories ############################ | ||
with tf.variable_scope("all"): | ||
make_directories(checkpoint_dir,save_dir) | ||
|
||
####### inputs ################### | ||
real_models = tf.placeholder(tf.float32, [args.batchsize, output_size, output_size, output_size] , name='real_models') | ||
z = tf.random_normal((args.batchsize, 200), 0, 1) | ||
a = tf.Print(z, [z], message="This is a: ") | ||
########## network computations ####################### | ||
|
||
net_g , G_train = generator_32(z, is_train=True, reuse = False, sig= True, batch_size=args.batchsize) | ||
|
||
if args.mini_batch_discimination: | ||
dis = discriminator_batch_discrimination | ||
else: | ||
dis = discriminator | ||
net_d , D_fake = dis(G_train, output_size, batch_size= args.batchsize, sig = True, is_train = True, reuse = False) | ||
net_d2, D_legit = dis(real_models, output_size, batch_size= args.batchsize, sig = True, is_train= True, reuse = True) | ||
net_d2, D_eval = dis(real_models, output_size, batch_size= args.batchsize, sig = True, is_train= False, reuse = True) # this is for desciding weather to train the discriminator | ||
|
||
|
||
########### Loss calculations ######################### | ||
d_loss = -tf.reduce_mean(tf.log(D_legit) + tf.log(1. - D_fake)) | ||
g_loss = -tf.reduce_mean(tf.log(D_fake)) | ||
|
||
|
||
############ Optimization ############# | ||
g_vars = net_g.all_params | ||
d_vars = net_d.all_params | ||
|
||
if args.mini_batch_discimination: | ||
W_var = [var for var in tf.trainable_variables() if var.name =='w_var:0'] | ||
d_vars += W_var | ||
|
||
net_g.print_params(False) | ||
net_d.print_params(False) | ||
|
||
d_optim = tf.train.AdamOptimizer(args.discriminator_learning_rate, beta1=0.5).minimize(d_loss, var_list=d_vars) | ||
g_optim = tf.train.AdamOptimizer(args.genorator_learning_rate, beta1=0.5).minimize(g_loss, var_list=g_vars) | ||
|
||
|
||
####### Training ################ | ||
config = tf.ConfigProto() | ||
config.gpu_options.allow_growth = True | ||
sess=tf.Session() | ||
tl.ops.set_gpu_fraction(sess=sess, gpu_fraction=0.998) | ||
sess.run(tf.global_variables_initializer()) | ||
|
||
if args.load: | ||
load_networks(checkpoint_dir, sess, net_g, net_d, epoch = args.load_epoch) | ||
|
||
files,iter_counter = grab_files(args.data) | ||
Train_Dis = True | ||
if len(args.load_epoch)>1: | ||
start = int(args.load_epoch) | ||
else: | ||
start = 0 | ||
for epoch in range(start, args.epochs): | ||
random.shuffle(files) | ||
for idx in xrange(0, len(files)/args.batchsize): | ||
file_batch = files[idx*args.batchsize:(idx+1)*args.batchsize] | ||
models, start_time = make_inputs(file_batch) | ||
#training the discriminator and the VAE's encoder | ||
if Train_Dis: | ||
errD,_,ones = sess.run([d_loss, d_optim, D_legit] ,feed_dict={real_models: models}) | ||
else: | ||
ones = sess.run([d_loss] ,feed_dict={real_models: models}) | ||
errG,_,zeros,objects = sess.run([g_loss, g_optim, D_fake, G_train], feed_dict={}) | ||
Train_Dis = (cal_acc(zeros,ones)<0.95)# only train discriminator at certain level of accuracy | ||
|
||
print("Epoch: [%2d/%2d] [%4d/%4d] time: %4.4f, d_loss: %.8f, g_loss: %.8f" \ | ||
% (epoch, args.epochs, idx, len(files)/args.batchsize , time.time() - start_time, errD, errG)) | ||
#saving the model | ||
if np.mod(epoch, args.save) == 0: | ||
save_networks(checkpoint_dir,sess, net_g, net_d, epoch) | ||
#saving generated objects | ||
if np.mod(epoch, args.sample ) == 0: | ||
save_voxels(save_dir,objects, epoch ) | ||
|
||
|
||
|
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,131 @@ | ||
import tensorflow as tf | ||
import keras.backend as K | ||
import os | ||
import sys | ||
import tensorlayer as tl | ||
import numpy as np | ||
from tensorlayer.layers import * | ||
import random | ||
from glob import glob | ||
import argparse | ||
import scripts | ||
from scripts.GANutils import * | ||
from scripts.models import * | ||
|
||
|
||
parser = argparse.ArgumentParser(description='3D-GAN implementation for 32*32*32 voxel output') | ||
parser.add_argument('-n','--name', default='Test', help='The name of the current experiment, this will be used to create folders and save models.') | ||
parser.add_argument('-d','--data', default='data/train/chair', help ='The location fo the object voxel models.' ) | ||
parser.add_argument('-e','--epochs', default=1500, help ='The number of epochs to run for.', type=int) | ||
parser.add_argument('-b','--batchsize', default=256, help ='The batch size.', type=int) | ||
parser.add_argument('-sample', default= 5, help='How often generated obejcts are sampled and saved.', type= int) | ||
parser.add_argument('-save', default= 5, help='How often the network models are saved.', type= int) | ||
parser.add_argument('-l', '--load', default= False, help='Indicates if a previously loaded model should be loaded.', action = 'store_true') | ||
parser.add_argument('-le', '--load_epoch', default= '', help='The epoch to number to be loaded from.', type=str) | ||
parser.add_argument('-graph', default= 5, help='How often the discriminator loss and the reconstruction loss graphs are saved.', type= int) | ||
args = parser.parse_args() | ||
|
||
checkpoint_dir = "checkpoint/" + args.name +'/' | ||
save_dir = "savepoint/" + args.name +'/' | ||
output_size = 32 | ||
|
||
|
||
with tf.variable_scope("all"): | ||
######### make directories ################ | ||
make_directories(checkpoint_dir,save_dir) | ||
|
||
####### inputs ################### | ||
z = tf.random_normal((args.batchsize, 200), 0, 1) | ||
real_models = tf.placeholder(tf.float32, [args.batchsize, output_size, output_size, output_size] , name='real_models') | ||
########## network computations ####################### | ||
|
||
|
||
|
||
#used for training genorator | ||
net_g, G_Fake = generator_32(z, is_train=True, reuse = False, sig= False, batch_size=args.batchsize) | ||
|
||
|
||
#used for training d on fake | ||
net_d, D_Fake = discriminator(G_Fake, output_size, batch_size= args.batchsize, improved = True ,is_train = True, reuse= False) | ||
#used for training d on real | ||
net_d2, D_Real = discriminator(real_models, output_size, batch_size= args.batchsize, improved = True ,is_train = True, reuse= True) | ||
|
||
########### Loss calculations ############ | ||
|
||
alpha = tf.random_uniform(shape=[args.batchsize,1] ,minval =0., maxval=1.) # here we calculate the gradient penalty | ||
difference = G_Fake - real_models | ||
inter = [] | ||
for i in range(args.batchsize): | ||
inter.append(difference[i] *alpha[i]) | ||
inter = tf.stack(inter) | ||
interpolates = real_models + inter | ||
gradients = tf.gradients(discriminator(interpolates, output_size, batch_size= args.batchsize, improved = True, is_train = False)[1],[interpolates])[0] | ||
slopes = tf.sqrt(tf.reduce_sum(tf.square(gradients),reduction_indices=[1])) | ||
gradient_penalty = tf.reduce_mean((slopes-1.)**2.) | ||
|
||
d_loss = -tf.reduce_mean(D_Real) + tf.reduce_mean(D_Fake) + 10.*gradient_penalty | ||
g_loss = -tf.reduce_mean(D_Fake) | ||
|
||
############ Optimization ############# | ||
|
||
g_vars = net_g.all_params # only updates the generator | ||
d_vars = net_d.all_params # only updates the discriminator | ||
|
||
|
||
net_g.print_params(False) | ||
net_d.print_params(False) | ||
|
||
d_optim = tf.train.AdamOptimizer( learning_rate = 1e-4, beta1=0.5, beta2=0.9).minimize(d_loss, var_list=d_vars) | ||
g_optim = tf.train.AdamOptimizer( learning_rate = 1e-4, beta1=0.5, beta2=0.9).minimize(g_loss, var_list=g_vars) | ||
|
||
|
||
####### Training ################ | ||
sess=tf.Session() | ||
tl.ops.set_gpu_fraction(sess=sess, gpu_fraction=0.998) | ||
sess.run(tf.global_variables_initializer()) | ||
|
||
# load checkpoints | ||
if args.load: | ||
load_networks(checkpoint_dir, sess, net_g, net_d, epoch = args.load_epoch) | ||
track_d_loss_iter, track_d_loss,_ = load_values(save_dir) | ||
else: | ||
track_d_loss_iter, track_d_loss, iter_counter = [],[],0 | ||
|
||
iter_counter = iter_counter - (iter_counter %5) | ||
files,_ = grab_files(args.data) | ||
#training starts here | ||
for epoch in range(args.epochs): | ||
random.shuffle(files) | ||
for idx in xrange(0, len(files)/args.batchsize): | ||
file_batch = files[idx*args.batchsize:(idx+1)*args.batchsize] | ||
models, start_time = make_inputs(file_batch) | ||
|
||
# updates the discriminator | ||
errD,_= sess.run([d_loss, d_optim] , feed_dict={ real_models: models }) | ||
track_d_loss.append(-errD) | ||
track_d_loss_iter.append(iter_counter) | ||
|
||
# update the generator | ||
if iter_counter % 5 ==0 : | ||
errG, _, objects= sess.run([g_loss, g_optim, G_Fake], feed_dict={}) | ||
|
||
print("Epoch: [%2d/%2d] [%4d/%4d] time: %4.4f, d_loss: %.8f, g_loss: %.8f" % (epoch, args.epochs, idx, len(files)/args.batchsize, time.time() - start_time, errD, errG)) | ||
sys.stdout.flush() | ||
|
||
iter_counter += 1 | ||
|
||
#saving generated objects | ||
if np.mod(epoch, args.sample ) == 0: | ||
save_voxels(save_dir,objects, epoch) | ||
#saving the model | ||
if np.mod(epoch, args.save) == 0: | ||
save_networks(checkpoint_dir,sess, net_g, net_d, epoch) | ||
|
||
|
||
#saving learning info | ||
if np.mod(epoch, args.graph) == 0: | ||
render_graphs(save_dir,epoch, track_d_loss_iter, track_d_loss) #this will only work after a 50 iterations to allows for proper averating | ||
save_values(save_dir,track_d_loss_iter, track_d_loss) # same here but for 300 | ||
|
||
|
||
|
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,9 @@ | ||
#!/bin/sh | ||
|
||
# get data | ||
|
||
wget http:https://3dshapenets.cs.princeton.edu/3DShapeNetsCode.zip | ||
unzip 3DShapeNetsCode | ||
|
||
# convert to our format | ||
python convert_shapenet10.py 3DShapeNets |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,6 @@ | ||
This is the directory for 3D generation. | ||
|
||
Two models are available, the first is my implementation of the 3D-GAN paper, found here: http:https://3dgan.csail.mit.edu/. The second is my own model, called 3D-IWGAN, released for my paper: . The dataset you can use to train is the modelNet10 dataset which can be downloaded and converted using the Make_Data.sh script. This will not take long. You can train by calling python 32-3D-IWGan.py [-h] [-n NAME] [-d DATA] [-e EPOCHS] [-b BATCHSIZE][-sample SAMPLE] [-save SAVE] [-l] [-le LOAD_EPOCH] [-graph GRAPH]. There is a description for all of these parameters by applying the '-h' parameter, though none are necessary and it will start training without then on the chair class with 12 orientations. | ||
To evaluate the output you can use the visualize.py script. To use this just call python visualize.py VOXEL_FILE, where voxel file is the file created during training and saved to savepoint/NAME/EPOCH.npy. For the 3D-IWGAN model, graphs will also be created which allow you to track the discriminator's loss, which will at fist raise rapidly but then begin to decrease, and this decreasing loss can be used to track convergence. To train on all of the training classes call 32-3D-IWGan.py -d 'data/train/*'. | ||
Let me know if you have any issues at [email protected] | ||
|
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,68 @@ | ||
import numpy as np | ||
import sys | ||
import os | ||
import scipy.io | ||
from path import Path | ||
from tqdm import tqdm | ||
|
||
if sys.argv[-1] == '-v': # this will allow you to visualize the models as they are made, more of a sanity check | ||
import mayavi.mlab | ||
import matplotlib.pyplot as plt | ||
from scipy import ndimage | ||
from mpl_toolkits.mplot3d import Axes3D | ||
|
||
|
||
instances = {} | ||
class_id_to_name = { | ||
"1": "bathtub", | ||
"2": "bed", | ||
"3": "chair", | ||
"4": "desk", | ||
"5": "dresser", | ||
"6": "monitor", | ||
"7": "night_stand", | ||
"8": "sofa", | ||
"9": "table", | ||
"10": "toilet" | ||
} | ||
class_name_to_id = { v : k for k, v in class_id_to_name.items() } | ||
class_names = set(class_id_to_name.values()) | ||
|
||
|
||
if not os.path.exists('data/train/'): | ||
os.makedirs('data/train/') | ||
if not os.path.exists('data/test/'): | ||
os.makedirs('data/test/') | ||
base_dir = Path(sys.argv[1]).expand() | ||
|
||
for fname in tqdm(sorted(base_dir.walkfiles('*.mat'))): | ||
if fname.endswith('test_feature.mat') or fname.endswith('train_feature.mat'): | ||
continue | ||
elts = fname.splitall() | ||
info = Path(elts[-1]).stripext().split('_') | ||
if len(info)<3: continue | ||
if info[0] == 'discriminative' or info[0] == 'generative' : continue | ||
instance = info[1] | ||
rot = int(info[2]) | ||
split = elts[-2] | ||
classname = elts[-4].strip() | ||
if classname in class_names: | ||
dest = 'data/'+split+'/' + classname + '/' | ||
if not os.path.exists(dest): | ||
os.makedirs(dest) | ||
arr = scipy.io.loadmat(fname)['instance'].astype(np.uint8) | ||
matrix = np.zeros((32,)*3, dtype=np.uint8) | ||
matrix[1:-1,1:-1,1:-1] = arr | ||
if sys.argv[-1] == '-v': | ||
xx, yy, zz = np.where(matrix>= 0.3) | ||
mayavi.mlab.points3d(xx, yy, zz, | ||
|
||
color=(.1, 0, 1), | ||
scale_factor=1) | ||
|
||
mayavi.mlab.show() | ||
# saves the models by instance name, and then rotation | ||
np.save(dest + instance + '_' + str(rot) , matrix) | ||
|
||
|
||
|
Oops, something went wrong.