fac_sortem.py 4.22 KB
Newer Older
Jonathan Juhl's avatar
Jonathan Juhl committed
1
2
3
import tensorflow as tf
import numpy as np

Jonathan Juhl's avatar
Jonathan Juhl committed
4
5
6
7
8
9
10
11
from trainer_sortem import Trainer
from mrc_loader_sortem import mrc_loader
class GAN_NERF():

    def __init__(self,args):
        self.args = args
        self.predict_steps = int(np.ceil(args['number_particles']/(args['num_gpus']*args['p_batch_size'])))
        steps = 0
12
13
14
15
16
17
18
19
20
21
22
23
        stages_list =  []
        l = [32,64,128,256]
        c = [args['s_1'],args['s_2'],args['s_3'],args['s_4']]
        for en,i in enumerate(l):
            if   i <= args['size']:
                steps+=c[en]
                stages_list.append(en*np.ones(c[en]))
            else:
                stages_list.append(np.concatenate([en*np.ones(c[en]),4*np.ones(self.args['top_off'])],axis=0))
            break
        self.stages = np.concatenate(stages_list)

Jonathan Juhl's avatar
Jonathan Juhl committed
24
25
26
27
28
29
        self.max_steps = steps
 
        gpu_list = []
        if args['gpu_list'] is None:
            args['gpu_list']  = []
        if  args['num_gpus'] > 1:
Jonathan Juhl's avatar
test    
Jonathan Juhl committed
30
            
Jonathan Juhl's avatar
Jonathan Juhl committed
31
32
33
34
35
36
37
            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
38
                
Jonathan Juhl's avatar
Jonathan Juhl committed
39
            else:
Jonathan Juhl's avatar
Jonathan Juhl committed
40
41
42
43
44
                    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
45
            else:
Jonathan Juhl's avatar
Jonathan Juhl committed
46
                gpu_list.append('GPU:%i' %args['gpu_list'][0])
Jonathan Juhl's avatar
Jonathan Juhl committed
47

Jonathan Juhl's avatar
Jonathan Juhl committed
48
49
        generator = mrc_loader(args).generate()
        predict_generator = mrc_loader(args).pred_generate() 
Jonathan Juhl's avatar
Jonathan Juhl committed
50

Jonathan Juhl's avatar
Jonathan Juhl committed
51
52
53
54
        if args['num_gpus'] > 1:
            strategy  = tf.distribute.MirroredStrategy(devices= gpu_list )
            self.generator = iter(strategy.experimental_distribute_dataset( generator))
            self.generator_pred  = iter(strategy.experimental_distribute_dataset(  predict_generator ))
Jonathan Juhl's avatar
Jonathan Juhl committed
55
     
Jonathan Juhl's avatar
Jonathan Juhl committed
56
        else:   
Jonathan Juhl's avatar
Jonathan Juhl committed
57
       
Jonathan Juhl's avatar
Jonathan Juhl committed
58
59
60
            strategy = tf.distribute.OneDeviceStrategy(device=gpu_list[0])
            self.generator = iter(strategy.experimental_distribute_dataset(  generator ))
            self.generator_pred = iter(strategy.experimental_distribute_dataset( predict_generator ))
Jonathan Juhl's avatar
Jonathan Juhl committed
61
     
Jonathan Juhl's avatar
Jonathan Juhl committed
62

Jonathan Juhl's avatar
Jonathan Juhl committed
63
        args['strategy'] = strategy
64
        args['max_steps'] = self.max_steps
Jonathan Juhl's avatar
Jonathan Juhl committed
65
        self.trainer = Trainer(args)
66
        #self.train()
Jonathan Juhl's avatar
Jonathan Juhl committed
67
68
69
70
71
        self.predict()
    def train(self):
        print('Begin training: ', '-' * 60)
        current_step = self.trainer.step_variable
   
72
73
74
            
        for i in range(int(current_step)):
            data = next(self.generator)
Jonathan Juhl's avatar
Jonathan Juhl committed
75
        for i in range(self.max_steps -current_step): # continue where you came from
76
77
78
79
80
81
82
83
84
85
86
87
88
            print("data step %i" %i )
            data = next(self.generator)
           
            if self.args['num_gpus'] == 1:
                    self.trainer.single_device_train_step(data,self.stages[i])
            else:
                    self.trainer.distributed_training_step(data,self.stages[i])
            if  (i % self.args['vi']) == 0:
                    self.trainer.write_summaries()
            if (i % self.args['save_model']) == 0:
                    self.trainer.save_checkpoint()
                    self.trainer.save_best_model()
            if i % self.args['movie_int']:
Jonathan Juhl's avatar
Jonathan Juhl committed
89
                if self.args['num_gpus'] == 1:
90
91
92
                    self.single_device_model_maker.model_maker()
                else:
                    self.multi_device_model_maker.model_maker()
93

Jonathan Juhl's avatar
Jonathan Juhl committed
94
    def predict(self):
Jonathan Juhl's avatar
Jonathan Juhl committed
95

96
97
        self.trainer.load_best_model()            
        #self.trainer.model_maker() 
Jonathan Juhl's avatar
Jonathan Juhl committed
98
99
        clusters = []
        for i in range(self.predict_steps):
100
101
102
103
            image = next(self.generator_pred )
    
            if self.args['num_gpus'] == 1:
               current = self.trainer.single_device_prediction_step(i)
Jonathan Juhl's avatar
Jonathan Juhl committed
104
            else:
105
106
107
108
               current = self.trainer.dis_predict_step(i)
            
            clusters.append(current)
        np.save(join(self.args['results'],'class_labels.npy'),np.asarray(clusters).flatten()[:self.args['number_particles']])
109
110


Jonathan Juhl's avatar
Jonathan Juhl committed
111