fac_sortem.py

import tensorflow as tf
import numpy as np

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
        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)

        self.max_steps = steps
 
        gpu_list = []
        if args['gpu_list'] is None:
            args['gpu_list']  = []
        if  args['num_gpus'] > 1:
            
            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)
                
            else:
                    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')
            else:
                gpu_list.append('GPU:%i' %args['gpu_list'][0])

        generator = mrc_loader(args).generate()
        predict_generator = mrc_loader(args).pred_generate() 

        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 ))
     
        else:   
       
            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 ))
     

        args['strategy'] = strategy
        args['max_steps'] = self.max_steps
        self.trainer = Trainer(args)
        #self.train()
        self.predict()
    def train(self):
        print('Begin training: ', '-' * 60)
        current_step = self.trainer.step_variable
   
            
        for i in range(int(current_step)):
            data = next(self.generator)
        for i in range(self.max_steps -current_step): # continue where you came from
            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']:
                if self.args['num_gpus'] == 1:
                    self.single_device_model_maker.model_maker()
                else:
                    self.multi_device_model_maker.model_maker()

    def predict(self):

        self.trainer.load_best_model()            
        #self.trainer.model_maker() 
        clusters = []
        for i in range(self.predict_steps):
            image = next(self.generator_pred )
    
            if self.args['num_gpus'] == 1:
               current = self.trainer.single_device_prediction_step(i)
            else:
               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']])