Commit 9ec83b9f authored by Jonathan Juhl's avatar Jonathan Juhl
Browse files

corrected log file

parent ecfe932b
import tensorflow as tf import tensorflow as tf
import numpy as np import numpy as np
import umap import umap
from os.path import join
from trainer_sortem import Trainer from trainer_sortem import Trainer
from mrc_loader_sortem import mrc_loader from mrc_loader_sortem import mrc_loader
class GAN_NERF(): class GAN_NERF():
...@@ -10,7 +10,7 @@ class GAN_NERF(): ...@@ -10,7 +10,7 @@ class GAN_NERF():
self.args = args self.args = args
dic = {32:1,64:2,128:3,256:4,512:5} dic = {32:1,64:2,128:3,256:4,512:5}
self.predict_steps = int(np.ceil(args['number_particles']/(args['num_gpus']*args['batch_size']))) self.predict_steps = int(np.ceil(args['number_particles']/(args['num_gpus']*args['batch_size'])))
self.dic = dic
l = np.asarray([32,64,128,256,512]) l = np.asarray([32,64,128,256,512])
self.args['resize'] = l[np.argmin(np.abs(l-self.args['size']))] self.args['resize'] = l[np.argmin(np.abs(l-self.args['size']))]
l_list = [] l_list = []
...@@ -52,17 +52,20 @@ class GAN_NERF(): ...@@ -52,17 +52,20 @@ class GAN_NERF():
predict_generator = mrc_loader(args).pred_generate() predict_generator = mrc_loader(args).pred_generate()
output_generator = mrc_loader(args).pred_generate()
if args['num_gpus'] > 1: if args['num_gpus'] > 1:
strategy = tf.distribute.MirroredStrategy(devices= gpu_list ) strategy = tf.distribute.MirroredStrategy(devices= gpu_list )
self.generator = strategy.experimental_distribute_dataset( generator) self.generator = strategy.experimental_distribute_dataset( generator)
self.generator_pred = strategy.experimental_distribute_dataset( predict_generator ) self.generator_pred = strategy.experimental_distribute_dataset( predict_generator )
self.output_generator = strategy.experimental_distribute_dataset( output_generator )
else: else:
strategy = tf.distribute.OneDeviceStrategy(device=gpu_list[0]) strategy = tf.distribute.OneDeviceStrategy(device=gpu_list[0])
self.generator = strategy.experimental_distribute_dataset( generator ) self.generator = strategy.experimental_distribute_dataset( generator )
self.generator_pred = strategy.experimental_distribute_dataset( predict_generator ) self.generator_pred = strategy.experimental_distribute_dataset( predict_generator )
self.output_generator = strategy.experimental_distribute_dataset( output_generator )
args['strategy'] = strategy args['strategy'] = strategy
self.trainer = Trainer(args) self.trainer = Trainer(args)
self.train() self.train()
...@@ -72,7 +75,7 @@ class GAN_NERF(): ...@@ -72,7 +75,7 @@ class GAN_NERF():
print('Begin training: ', '-' * 60) print('Begin training: ', '-' * 60)
current_step = self.trainer.step_variable current_step = self.trainer.step_variable
gen = iter(self.generator) gen = iter(self.generator)
pred = iter(self.generator_pred)
for i in range(int(current_step)): for i in range(int(current_step)):
# this starts the data recording at where it left off # this starts the data recording at where it left off
# this is to prevent when continuation of training the model does not use the same data # this is to prevent when continuation of training the model does not use the same data
...@@ -101,27 +104,52 @@ class GAN_NERF(): ...@@ -101,27 +104,52 @@ class GAN_NERF():
self.trainer.distributed_training_step(params) self.trainer.distributed_training_step(params)
if (i % self.args['record']) == 0: if (i % self.args['record']) == 0:
if self.args['num_gpus'] == 1: if self.args['num_gpus'] == 1:
self.single_device_model_maker.model_maker() features = []
current_shape = params['image'].numpy().shape[1]
for kk in range(int(np.ceil(self.args['umap_t_size']/self.args['batch_size']))):
data = next(pred)
features.append(self.args['strategy'].run(self.trainer.get_features,
args=(data,params['alpha'],self.trainer.Encoder[int(params['index'])],current_shape)))
self.trainer.write_summaries(features)
else: else:
self.multi_device_model_maker.model_maker() data = next(pred)
self.trainer.write_summaries() features = []
for kk in range(int(np.ceil(self.args['umap_t_size']/(self.num_gpus*self.args['batch_size'])))):
features.append(self.args['strategy'].run(self.trainer.get_features,
args=(data,params['alpha'],self.trainer.Encoder[int(params['index'])],current_shape)).reduce())
self.trainer.write_summaries(features)
if (i % self.args['save_model']) == 0: if (i % self.args['save_model']) == 0:
self.trainer.save_checkpoint() self.trainer.save_checkpoint()
self.trainer.save_best_model()
def over_cluster(self):
self.trainer.load_best_model()
trainer.sparse_water_sheed_algorithm()
def predict(self): def predict(self):
pred = iter(self.output_generator)
output_vectors = []
if not isfile(join(self.args['results'],'final_featur_vectors.npy')):
if self.args['num_gpus'] > 1:
for kk in range(int(np.ceil(self.args['depth']/self.args['num_gpus']*self.args['batch_size']))):
data = next(pred)
output_vectors.append(self.args['strategy'].run(self.trainer.get_features,args=(data,params['alpha'],self.trainer.Encoder[self.dic[self.args['resize']]],current_shape)).reduce())
else:
for kk in range(int(np.ceil(self.args['depth']/self.args['batch_size']))):
data = next(pred)
output_vectors.append(self.args['strategy'].run(self.trainer.get_features,args=(data,params['alpha'],self.trainer.Encoder[self.dic[self.args['resize']]],current_shape)))
np.save(join(self.args['results'],'final_featur_vectors.npy'))
labels,umap_output,collect_centers = pred_umap(args,feature_vector)
if not isfile(join(self.args['results'],'final_labels.npy')):
np.save(join(self.args['results'],'final_labels.npy'))
np.save(join(self.args['results'],'final_umap_output.npy'))
np.save(join(self.args['results'],'final_collect_centers.npy'))
self.trainer.load_best_model()
#self.trainer.model_maker()
bools = isfile(join(self.args['results'],'over_cluster.npy'))
if bools:
labels = np.load(join(self.args['results'],'over_cluster.npy'))
clusters = []
for i in range(self.predict_steps): for i in range(self.predict_steps):
image = next(self.generator_pred ) image = next(self.generator_pred )
......
...@@ -33,7 +33,7 @@ class ED_Maker(Model): ...@@ -33,7 +33,7 @@ class ED_Maker(Model):
self.final_layer = final_layer self.final_layer = final_layer
self.fromRGB.reverse() self.fromRGB.reverse()
self.start = self.img_size_to_layer[shape] self.start = self.img_size_to_layer[shape]
self.flatten = Flatten()
self.lays = layers self.lays = layers
def call(self,input,alpha): def call(self,input,alpha):
...@@ -51,7 +51,7 @@ class ED_Maker(Model): ...@@ -51,7 +51,7 @@ class ED_Maker(Model):
x = layer(x) x = layer(x)
out = tf.squeeze(self.final_layer(x)) out = self.flatten(self.final_layer(x))
...@@ -61,7 +61,8 @@ class ED_Maker(Model): ...@@ -61,7 +61,8 @@ class ED_Maker(Model):
v = [] v = []
for i in self.lays[:self.start ]: for i in self.lays[:self.start ]:
v+= i.trainable_variables v+= i.trainable_variables
return v+self.fromRGB[self.start-1].trainable_variables+self.fromRGB[ int(self.start) ].trainable_variables k = self.final_layer.trainable_variables
return v+k+self.fromRGB[self.start-1].trainable_variables+self.fromRGB[ int(self.start) ].trainable_variables
class ResidualConvBlock(Layer): class ResidualConvBlock(Layer):
def __init__(self,inplanes, planes, kernel_size=3, stride=1, downsample=False, groups=1): def __init__(self,inplanes, planes, kernel_size=3, stride=1, downsample=False, groups=1):
...@@ -167,7 +168,7 @@ class Ray_maker(Layer): ...@@ -167,7 +168,7 @@ class Ray_maker(Layer):
s = tf.shape(coordinates)[1] s = tf.shape(coordinates)[1]
powered = tf.expand_dims(tf.pow(tf.cast(2.0,dtype=coordinates.dtype),tf.cast(tf.range(self.init_L),dtype=coordinates.dtype)),axis=0) powered = tf.expand_dims(tf.pow(tf.cast(2.0,dtype=coordinates.dtype),tf.cast(tf.range(self.init_L),dtype=coordinates.dtype)),axis=0)
out = tf.matmul(coordinates,powered,coordinates.dtype) out = tf.matmul(coordinates,powered,transpose_a=True)
out = tf.reshape(out,[-1,s,self.init_L*3]) out = tf.reshape(out,[-1,s,self.init_L*3])
...@@ -212,9 +213,9 @@ class Ray_maker(Layer): ...@@ -212,9 +213,9 @@ class Ray_maker(Layer):
a,b = tf.split(z,2,axis=-1) a,b = tf.split(z,2,axis=-1)
a = tf.squeeze(a) a = tf.squeeze(a)
b = tf.squeeze(b) b = tf.squeeze(b)
fourie = tf.complex(tf.cast(a,tf.float32),tf.cast(b,tf.float32))
return tf.squeeze(fourie) return a,b
class Noise_Maker(Model): class Noise_Maker(Model):
......
...@@ -64,10 +64,8 @@ class mrc_loader: ...@@ -64,10 +64,8 @@ class mrc_loader:
image = tf.io.decode_raw(ins,tf.float32) image = tf.io.decode_raw(ins,tf.float32)
image = tf.reshape(image,[self.df_keys['size'],self.df_keys['size'],1]) image = tf.reshape(image,[self.df_keys['size'],self.df_keys['size'],1])
image = tf.image.per_image_standardization(image) image = tf.image.per_image_standardization(image)
minima = tf.reduce_min(image) image = image-tf.reduce_min(image)
tmp = image+minima return image/tf.reduce_max(image)
image = tmp/tf.reduce_max(tmp)
return image
data = tf.data.FixedLengthRecordDataset(self.df_keys['mrc_paths'],self.df_keys['bpr'],num_parallel_reads=self.df_keys['num_cpus'], header_bytes=1024).map(pred_image,self.df_keys['num_cpus']).batch(self.df_keys['batch_size']).repeat() data = tf.data.FixedLengthRecordDataset(self.df_keys['mrc_paths'],self.df_keys['bpr'],num_parallel_reads=self.df_keys['num_cpus'], header_bytes=1024).map(pred_image,self.df_keys['num_cpus']).batch(self.df_keys['batch_size']).repeat()
return data return data
...@@ -75,3 +73,25 @@ class mrc_loader: ...@@ -75,3 +73,25 @@ class mrc_loader:
# def generator_model(self,input_vectors,): # def generator_model(self,input_vectors,):
class Grid_Maker:
def __init__(self,kwargs,means):
self.size = kwargs['size']
x = tf.linspace(0.0,0.5,int(self.size/2)+1)
y = tf.linspace(-0.5,0.5,self.size)
z = tf.linspace(-0.5,0.5,self.size)
X,Y,Z = tf.meshgrid(x,y,z)
self.X = X
self.Y = Y
self.Z = Z
self.mean_size = means.shape[0]
def grid_generator(self):
x_slice = tf.from_tensor_slices(self.X)
y_slice = tf.from_tensor_slices(self.Y)
z_slice = tf.from_tensor_slices(self.Z)
return tf.data.Dataset.zip(x_slice,y_slice,z_slice).repeat(self.mean_size).batch(self.kwargs['m_batch_size'])
This diff is collapsed.
...@@ -7,11 +7,52 @@ from os.path import join ...@@ -7,11 +7,52 @@ from os.path import join
from tensorflow.keras import Model from tensorflow.keras import Model
# Clustering Analysis via Deep Generative Models With Mixture Models # Clustering Analysis via Deep Generative Models With Mixture Models
from tensorflow.keras.layers import Layer from tensorflow.keras.layers import Layer
import umap
import hdbscan
from sklearn.cluster import KMeans
def get_parameters(self,normed_rotation):
normed_rotation,_ = tf.linalg.normalize(normed_rotation,axis=1)
translate_x = 2*translation[:,0]*self.ny
translate_y = 2*translation[:,1]*self.nx
tilt = tf.math.atan2(normed_rotation[:,2,1],normed_rotation[:,2,2])*180/(np.pi)
psi = tf.math.atan2(-normed_rotation[:,2,0],tf.sqrt(tf.pow(normed_rotation[:,2,1],2)+tf.pow(normed_rotation[:,2,2],2)))*180/(np.pi)
euler = tf.math.atan2(normed_rotation[:,1,0],normed_rotation[:,0,0])*180/(np.pi)
inplane = tf.reshape(tf.stack([tf.cos(euler),-tf.sin(euler),tf.sin(euler),tf.cos(euler)],axis=-1),[-1,2,2])
return tilt,psi,inplane
class Make_2D_Transform(tf.keras.layers.Layer):
def __init__(self,width):
super(Make_2D_Transform, self).__init__()
self.width = width
fourier = np.fft.fftfreq(width)
x = tf.linspace(0.0,float(width)-1,width)
self.F_X,self.X = np.meshgrid(x,fourier)
def call(self,real,imaginary):
s_r = tf.reverse(real,axis=[1])
s_i = tf.reverse(imaginary,axis=[1])
s_r = tf.concat([tf.expand_dims(s_r[:,0,:],axis=1),tf.reverse(s_r[:,1:,:],axis=[1])],axis=1)
s_im = tf.concat([tf.expand_dims(s_i[:,0,:],axis=1),tf.reverse(s_i[:,1:,:],axis=[1])],axis=1)
a = tf.concat([real,s_r[:,:,1:-1]],axis=2)
b = tf.concat([imaginary,-s_im[:,:,1:-1]],axis=2)
A = tf.cast(tf.cos(2*np.pi*self.F_X*self.X),tf.float32)/self.width
B = tf.cast(tf.sin(2*np.pi*self.F_X*self.X),tf.float32)/self.width
a_1 = tf.matmul(A,a)-tf.matmul(B,b)
b_1 = tf.matmul(B,a)+tf.matmul(A,b)
a_out = tf.matmul(a_1,A)-tf.matmul(b_1,B)
return a_out
class DiversityLoss(tf.keras.losses.Loss): class DiversityLoss(tf.keras.losses.Loss):
def call(self,z,x): def call(self,z,x):
nominator = (tf.norm(tf.math.roll(x,1,axis=0)-x)) nominator = (tf.norm(tf.roll(x,1,axis=0)-x))
denominator = (tf.norm(tf.math.roll(z,1,axis=0)-z)) denominator = (tf.norm(tf.roll(z,1,axis=0)-z))
s = -tf.math.minimum(nominator/denominator,10**(-4)) s = -tf.math.minimum(nominator/denominator,10**(-4))
return s return s
class AverageBlurPooling2D(Layer): class AverageBlurPooling2D(Layer):
...@@ -105,12 +146,12 @@ class Eval_CTF(tf.keras.layers.Layer): ...@@ -105,12 +146,12 @@ class Eval_CTF(tf.keras.layers.Layer):
return tf.cast(ctf,tf.complex64) return tf.cast(ctf,tf.complex64)
def gradient_penalty(self, discriminate,x_real, x_gen): def gradient_penalty(discriminate,x_real, x_gen,alpha):
epsilon = tf.random.uniform([x.shape[0], 1, 1, 1], 0.0, 1.0) epsilon = tf.random.uniform([x_real.shape[0], 1, 1, 1], 0.0, 1.0)
x_hat = epsilon *x_real + (1 - epsilon) * x_gen x_hat = epsilon *x_real + (1 - epsilon) * x_gen
with tf.GradientTape() as t: with tf.GradientTape() as t:
t.watch(x_hat) t.watch(x_hat)
d_hat = discriminate(x_hat) d_hat = discriminate(x_hat,alpha)
gradients = t.gradient(d_hat, x_hat) gradients = t.gradient(d_hat, x_hat)
ddx = tf.sqrt(tf.reduce_sum(gradients ** 2, axis=[1, 2])) ddx = tf.sqrt(tf.reduce_sum(gradients ** 2, axis=[1, 2]))
d_regularizer = tf.reduce_mean((ddx - 1.0) ** 2) d_regularizer = tf.reduce_mean((ddx - 1.0) ** 2)
...@@ -118,37 +159,39 @@ def gradient_penalty(self, discriminate,x_real, x_gen): ...@@ -118,37 +159,39 @@ def gradient_penalty(self, discriminate,x_real, x_gen):
class Discriminator_Loss(tf.keras.losses.Loss): class Discriminator_Loss(tf.keras.losses.Loss):
def init(self): def init(self):
self.loss_real = tf.keras.losses.BinaryCrossentropy(from_logits=True) super(Discriminator_Loss, self).__init__()
self.loss_fake = tf.keras.losses.BinaryCrossentropy(from_logits=True) def call(self,D_real, y_pred):
self.loss_tilde = tf.keras.losses.BinaryCrossentropy(from_logits=True) D_tilde,D_fake = y_pred
def call(self,x_real, D_real, y_pred):
D_fake,D_tilde,D_real,reg = y_pred
real_loss = self.loss_real(tf.ones_like(D_real),D_real)
tilde_loss = self.loss_fake(tf.zeros_like(D_tilde),D_tilde)
fake_loss = self.loss_tilde(tf.zeros_like(D_tilde),D_fake)
tilde_loss = tf.nn.sigmoid_cross_entropy_with_logits(tf.zeros_like(D_tilde),D_tilde)
return real_loss+tilde_loss+fake_loss+reg real_loss = tf.nn.sigmoid_cross_entropy_with_logits(tf.ones_like(D_real),D_real)
fake_loss = tf.nn.sigmoid_cross_entropy_with_logits(tf.zeros_like(D_tilde),D_fake)
gan_loss = real_loss+tilde_loss+fake_loss
return gan_loss
class Generator_Loss(tf.keras.losses.Loss): class Generator_Loss(tf.keras.losses.Loss):
def init(self): def init(self):
self.loss_mse = tf.keras.losses.MSE() super(Generator_Loss, self).__init__()
self.loss_tilde = tf.keras.losses.BinaryCrossentropy(from_logits=True) def call(self, D_real, y_pred):
self.loss_real = tf.keras.losses.BinaryCrossentropy(from_logits=True) tilde_image,real_image, D_tilde,D_fake = y_pred
def call(self, real_image, y_pred): real_loss = tf.nn.sigmoid_cross_entropy_with_logits(tf.ones_like(D_real),D_real)
D_tilde,D_fake,tilde_image = y_pred tilde_loss = tf.nn.sigmoid_cross_entropy_with_logits(tf.zeros_like(D_tilde),D_tilde)
fake_loss = tf.nn.sigmoid_cross_entropy_with_logits(tf.zeros_like(D_tilde),D_fake)
ll_loss = self.loss_mse(real_image,tilde_image) gan_loss = real_loss+tilde_loss+fake_loss
d_fake = self.loss_tilde(tf.ones_like(D_fake),D_fake)
d_real = self.loss_real(tf.ones_like(D_tilde),D_tilde)
return d_fake+d_real+0.1*ll_loss c = -0.5 * tf.math.log(2 * np.pi)
multiplier = 1.0 / (2.0 * 1)
tmp = tf.square(tilde_image - real_image)
tmp *= -multiplier
tmp += c
return 0.1 * tf.reduce_sum(tmp,axis=0) - gan_loss
class Encoder_Loss(tf.keras.losses.Loss): class Encoder_Loss(tf.keras.losses.Loss):
def init(self): def init(self):
self.loss_mse = tf.keras.losses.MSE() super(Encoder_Loss, self).__init__()
self.loss_tilde = tf.keras.losses.BinaryCrossentropy(from_logits=True)
self.loss_real = tf.keras.losses.BinaryCrossentropy(from_logits=True)
def kl(self,mu,log_var): def kl(self,mu,log_var):
return -0.5 * tf.reduce_sum(1 + log_var - tf.pow(mu, 2) - tf.exp(log_var)) return -0.5 * tf.reduce_sum(1 + log_var - tf.pow(mu, 2) - tf.exp(log_var))
...@@ -173,14 +216,16 @@ class Poisson_Measure(tf.keras.layers.Layer): ...@@ -173,14 +216,16 @@ class Poisson_Measure(tf.keras.layers.Layer):
super(Poisson_Measure, self).__init__() super(Poisson_Measure, self).__init__()
def call(self,input_image,noise_lam): def call(self,input_image,noise_lam):
# compute the approximate poisson noise and add to image # compute the approximate poisson noise and add to image
input_image = tf.expand_dims(input_image,axis=-1)
# rescale the image such that the square root is no negative # rescale the image such that the square root is no negative
mean = tf.constant([0.0],dtype=input_image.dtype);scale =tf.constant([2.0],dtype=input_image.dtype) mean = tf.constant([0.0],dtype=input_image.dtype);scale =tf.constant([2.0],dtype=input_image.dtype)
_noise_scale = tf.cast(tf.sqrt(1. / noise_lam),dtype=input_image.dtype) _noise_scale = tf.cast(tf.sqrt(1. / noise_lam),dtype=input_image.dtype)
eps = tf.random.uniform(shape=tf.shape(input_image),dtype=input_image.dtype) eps = tf.random.uniform(shape=tf.shape(input_image),dtype=input_image.dtype)
noise = (eps * _noise_scale * noise = (eps * _noise_scale *
tf.sqrt((tf.stop_gradient(input_image)- mean) / scale + 0.5)) * scale tf.sqrt((tf.stop_gradient(input_image)- mean) / scale + 0.5)) * scale
output = input_image + noise output = input_image + noise
return output return output
...@@ -196,12 +241,18 @@ def make_umap(args,feature_vectors): ...@@ -196,12 +241,18 @@ def make_umap(args,feature_vectors):
min_cluster_size=args['minimum_size'], min_cluster_size=args['minimum_size'],
).fit_predict(umap_output) ).fit_predict(umap_output)
index = np.unique(labels) index = np.unique(labels)
unique_labels = np.unique(labels)
unique_labels = unique_labels[np.greater(unique_labels,-1)]
collect_centers = [] collect_centers = []
for i in index:
bools = np.equal(i,labels) for i in unique_labels:
selected_features = np.take_along_axis(feature_vectors,bools,axis=0)
selected_features = umap_output[np.equal(i,labels)]
collect_centers.append(KMeans(n_clusters=args['frames']).fit(feature_vectors).cluster_centers_) collect_centers.append(KMeans(n_clusters=args['frames']).fit(feature_vectors).cluster_centers_)
return np.concatenate(collect_centers)
return labels,umap_output,np.concatenate(collect_centers)
class Make_Grids(tf.keras.layers.Layer): class Make_Grids(tf.keras.layers.Layer):
def __init__(self): def __init__(self):
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment