fac_sortem.py 7.9 KB
Newer Older
Jonathan Juhl's avatar
Jonathan Juhl committed
1
2
import tensorflow as tf
import numpy as np
Jonathan Juhl's avatar
Jonathan Juhl committed
3
import umap
Jonathan Juhl's avatar
Jonathan Juhl committed
4
from os.path import join
Jonathan Juhl's avatar
Jonathan Juhl committed
5
6
7
8
9
10
from trainer_sortem import Trainer
from mrc_loader_sortem import mrc_loader
class GAN_NERF():

    def __init__(self,args):
        self.args = args
Jonathan Juhl's avatar
Jonathan Juhl committed
11
12
        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'])))   
Jonathan Juhl's avatar
Jonathan Juhl committed
13
        self.dic = dic
Jonathan Juhl's avatar
Jonathan Juhl committed
14
15
16
17
18
19
20
21
22
        l = np.asarray([32,64,128,256,512])
        self.args['resize'] = l[np.argmin(np.abs(l-self.args['size']))]
        l_list = []
        alpha_list = []
        im_shape_list = []
        for index,steps in enumerate(args['steps'][:dic[self.args['resize']]]):
                l_list += [index]*steps
                alpha_list += [np.linspace(0.0,1.0,steps)]
                im_shape_list += [l[index]]*steps
23

Jonathan Juhl's avatar
Jonathan Juhl committed
24
25
26
27
28
        self.alphas = tf.constant(tf.concat(alpha_list,axis=0))
        self.index  = tf.constant(np.asarray(l_list))
        self.im_shape = tf.constant(np.asarray(im_shape_list))
        self.steps = self.alphas.shape[0]
     
Jonathan Juhl's avatar
Jonathan Juhl committed
29
30
31
32
        gpu_list = []
        if args['gpu_list'] is None:
            args['gpu_list']  = []
        if  args['num_gpus'] > 1:
Jonathan Juhl's avatar
test    
Jonathan Juhl committed
33
            
Jonathan Juhl's avatar
Jonathan Juhl committed
34
35
36
37
38
39
40
            if len(args['gpu_list']) == args['num_gpus'] or len(args['gpu_list']) == 0:
                if len(args['gpu_list']) == 0:
                    for i in range(args['num_gpus']):
                                gpu_list.append('GPU:%i' %i)
                else:
                    for i in args['gpu_list']:
                            gpu_list.append('GPU:%i' %i)
Jonathan Juhl's avatar
test    
Jonathan Juhl committed
41
                
Jonathan Juhl's avatar
Jonathan Juhl committed
42
            else:
Jonathan Juhl's avatar
Jonathan Juhl committed
43
44
45
46
47
                    print("the number of gpus specified is not equal to the number of gpu names. Either leave gpu names parameter unmodified or correct it.")
                    exit()
        else:
            if len(args['gpu_list']) == 0:
                gpu_list.append('GPU:0')
Jonathan Juhl's avatar
Jonathan Juhl committed
48
            else:
Jonathan Juhl's avatar
Jonathan Juhl committed
49
                gpu_list.append('GPU:%i' %args['gpu_list'][0])
Jonathan Juhl's avatar
Jonathan Juhl committed
50
51
   
        generator = mrc_loader(args).generate(self.alphas,self.index,self.im_shape)
Jonathan Juhl's avatar
Jonathan Juhl committed
52

Jonathan Juhl's avatar
Jonathan Juhl committed
53

Jonathan Juhl's avatar
Jonathan Juhl committed
54
        predict_generator = mrc_loader(args).pred_generate() 
Jonathan Juhl's avatar
Jonathan Juhl committed
55
        output_generator = mrc_loader(args).pred_generate() 
Jonathan Juhl's avatar
Jonathan Juhl committed
56
   
Jonathan Juhl's avatar
Jonathan Juhl committed
57
58
        if args['num_gpus'] > 1:
            strategy  = tf.distribute.MirroredStrategy(devices= gpu_list )
Jonathan Juhl's avatar
Jonathan Juhl committed
59
60
            self.generator = strategy.experimental_distribute_dataset( generator)
            self.generator_pred  = strategy.experimental_distribute_dataset(  predict_generator )
Jonathan Juhl's avatar
Jonathan Juhl committed
61
            self.output_generator  = strategy.experimental_distribute_dataset(  output_generator )
Jonathan Juhl's avatar
Jonathan Juhl committed
62
        else:   
Jonathan Juhl's avatar
Jonathan Juhl committed
63
       
Jonathan Juhl's avatar
Jonathan Juhl committed
64
            strategy = tf.distribute.OneDeviceStrategy(device=gpu_list[0])
Jonathan Juhl's avatar
Jonathan Juhl committed
65
66
            self.generator = strategy.experimental_distribute_dataset(  generator )
            self.generator_pred = strategy.experimental_distribute_dataset( predict_generator )
Jonathan Juhl's avatar
Jonathan Juhl committed
67
68
            self.output_generator  = strategy.experimental_distribute_dataset(  output_generator )
    
Jonathan Juhl's avatar
Jonathan Juhl committed
69
70
        args['strategy'] = strategy
        self.trainer = Trainer(args)
Jonathan Juhl's avatar
Jonathan Juhl committed
71
        self.train()
Jonathan Juhl's avatar
Jonathan Juhl committed
72
        self.predict()
Jonathan Juhl's avatar
Jonathan Juhl committed
73

Jonathan Juhl's avatar
Jonathan Juhl committed
74
75
76
    def train(self):
        print('Begin training: ', '-' * 60)
        current_step = self.trainer.step_variable
Jonathan Juhl's avatar
Jonathan Juhl committed
77
        gen = iter(self.generator)
Jonathan Juhl's avatar
Jonathan Juhl committed
78
        pred = iter(self.generator_pred)
Jonathan Juhl's avatar
Jonathan Juhl committed
79
80
81
82
83
84
        for i in range(int(current_step)): 
            # 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
            data = next(gen)
 
        for i in range(self.steps -current_step): # continues where we left off
85
            print("data step %i" %i )
Jonathan Juhl's avatar
Jonathan Juhl committed
86
87
88
            data = next(gen) 

            if self.args['ctf']:
Jonathan Juhl's avatar
Jonathan Juhl committed
89
                params = {'image':tf.expand_dims(tf.squeeze(data[0][0]),axis=-1),
Jonathan Juhl's avatar
Jonathan Juhl committed
90
91
                          'ctf':data[0][1],
                          'alpha':data[1],
Jonathan Juhl's avatar
Jonathan Juhl committed
92
                          'index':data[2],                     
Jonathan Juhl's avatar
Jonathan Juhl committed
93
94
                        }
            else:
Jonathan Juhl's avatar
Jonathan Juhl committed
95
                params = {'image':tf.expand_dims(tf.squeeze(data[0]),axis=-1),
Jonathan Juhl's avatar
Jonathan Juhl committed
96
97
                          'ctf': None,
                          'alpha':data[1],
Jonathan Juhl's avatar
Jonathan Juhl committed
98
                          'index':data[2],                         
Jonathan Juhl's avatar
Jonathan Juhl committed
99
100
                        }
          
101
            if self.args['num_gpus'] == 1:
Jonathan Juhl's avatar
Jonathan Juhl committed
102
                    self.trainer.single_device_train_step(params)
103
            else:
Jonathan Juhl's avatar
Jonathan Juhl committed
104
105
                    self.trainer.distributed_training_step(params)
            if  (i % self.args['record']) == 0:
Jonathan Juhl's avatar
Jonathan Juhl committed
106
                if self.args['num_gpus'] == 1:
Jonathan Juhl's avatar
Jonathan Juhl committed
107
108
109
110
111
112
113
114
115
116
                    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)
117
                else:
Jonathan Juhl's avatar
Jonathan Juhl committed
118
119
120
121
122
123
124
125
126
                    data = next(pred) 
                    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)
          
           
Jonathan Juhl's avatar
Jonathan Juhl committed
127
128
            if (i % self.args['save_model']) == 0:
                    self.trainer.save_checkpoint()
Jonathan Juhl's avatar
Jonathan Juhl committed
129
130
                  

131

Jonathan Juhl's avatar
Jonathan Juhl committed
132
    def predict(self):
Jonathan Juhl's avatar
Jonathan Juhl committed
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
        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'))
Jonathan Juhl's avatar
Jonathan Juhl committed
152

Jonathan Juhl's avatar
Jonathan Juhl committed
153
        for i in range(self.predict_steps):
154
            image = next(self.generator_pred )
Jonathan Juhl's avatar
Jonathan Juhl committed
155
          
156
            if self.args['num_gpus'] == 1:
Jonathan Juhl's avatar
Jonathan Juhl committed
157
               current = self.trainer.single_device_prediction_step(image)
Jonathan Juhl's avatar
Jonathan Juhl committed
158
            else:
Jonathan Juhl's avatar
Jonathan Juhl committed
159
160
161
162
163
164
165
166
167
168
                current = self.trainer.dis_predict_step(image)
                clusters.append(current.numpy())
        training_feat = np.asarray(clusters).flatten()[:self.args['umap_size'],:]
        predict_feat = np.asarray(clusters).flatten()[self.args['umap_size']:,:]
        training_set = umap.UMAP(neighbours = self.args['neighbours'])
        TD = training_set.fit_predict(training_feat)
        TD_after = training_set.fit_predict(predict_feat)
        information = np.concatenate([TD,TD_after],axis=0)
        labels = hdbscan.HDBSCAN(min_cluster_size=self.args['minimum_size'], min_samples=10).fet_predict(information)
        np.save(join(self.args['results'],'class_labels.npy'),labels[:self.args['number_particles']])
169
170


Jonathan Juhl's avatar
Jonathan Juhl committed
171