Skip to content
Snippets Groups Projects
GenLearner.py 8.54 KiB
Newer Older
  • Learn to ignore specific revisions
  • Max Ramsay King's avatar
    Max Ramsay King committed
    import torch
    
    
    import autoaug.child_networks as cn
    from autoaug.autoaugment_learners.AaLearner import AaLearner
    
    Max Ramsay King's avatar
    Max Ramsay King committed
    import random
    
    
    class Genetic_learner(AaLearner):
    
        def __init__(self, 
                    # search space settings
                    sp_num=5,
                    p_bins=11, 
                    m_bins=10, 
                    discrete_p_m=False,
                    exclude_method=[],
                    # child network settings
                    learning_rate=1e-1, 
                    max_epochs=float('inf'),
                    early_stop_num=20,
                    batch_size=8,
                    toy_size=1,
                    num_offspring=1, 
                    ):
    
            super().__init__(
                        sp_num=sp_num, 
                        p_bins=p_bins, 
                        m_bins=m_bins, 
                        discrete_p_m=discrete_p_m, 
                        batch_size=batch_size, 
                        toy_size=toy_size, 
                        learning_rate=learning_rate,
                        max_epochs=max_epochs,
                        early_stop_num=early_stop_num,
                        exclude_method=exclude_method
                        )
    
            self.bin_to_aug =  {}
            for idx, augmentation in enumerate(self.augmentation_space):
                bin_rep = '{0:b}'.format(idx)
                while len(bin_rep) < len('{0:b}'.format(len(self.augmentation_space))):
                    bin_rep = '0' + bin_rep
                self.bin_to_aug[bin_rep] = augmentation[0]
    
            self.just_augs = [x[0] for x in self.augmentation_space]
        
            self.mag_to_bin = {
                '0': "0000",
                '1': '0001',
                '2': '0010',
                '3': '0011',
                '4': '0100',
                '5': '0101',
                '6': '0110',
                '7': '0111',
                '8' : '1000',
                '9': '1001',
                '10': '1010',
            }
    
            self.prob_to_bin = {
                '0': "0000",
                '0.0' : '0000',
                '0.1': '0001',
                '0.2': '0010',
                '0.3': '0011',
                '0.4': '0100',
                '0.5': '0101',
                '0.6': '0110',
                '0.7': '0111',
                '0.8' : '1000',
                '0.9': '1001',
                '1.0': '1010',
                '1' : '1010',
            }
    
            self.bin_to_prob = dict((value, key) for key, value in self.prob_to_bin.items())
            self.bin_to_mag = dict((value, key) for key, value in self.mag_to_bin.items())
            self.aug_to_bin = dict((value, key) for key, value in self.bin_to_aug.items())
    
            self.num_offspring = num_offspring
    
    
        def gen_random_subpol(self):
    
    Max Ramsay King's avatar
    Max Ramsay King committed
            """
            Generates a random subpolicy using the reduced augmentation_space
    
            Returns
            --------
            subpolicy -> ((transformation, probability, magnitude), (trans., prob., mag.))
            """
    
    Max Ramsay King's avatar
    Max Ramsay King committed
            choose_items = [x[0] for x in self.augmentation_space]
            trans1 = str(random.choice(choose_items))
            trans2 = str(random.choice(choose_items))
            prob1 = float(random.randrange(0, 11, 1) / 10)
            prob2 = float(random.randrange(0, 11, 1) / 10)
    
            if self.aug_space_dict[trans1]:
                mag1 = int(random.randrange(0, 10, 1))
            else:
                mag1 = None
    
            if self.aug_space_dict[trans2]:
                mag2 = int(random.randrange(0, 10, 1))
            else:
                mag2 = None
    
            subpol = ((trans1, prob1, mag1), (trans2, prob2, mag2))
            return subpol
    
    
        def gen_random_policy(self):
    
    Max Ramsay King's avatar
    Max Ramsay King committed
            """
            Generates a random policy, consisting of sp_num subpolicies
    
            Returns
            ------------
            policy -> [subpolicy, subpolicy, ...]
            """
    
    Max Ramsay King's avatar
    Max Ramsay King committed
            pol = []
            for _ in range(self.sp_num):
                pol.append(self.gen_random_subpol())
            return pol
    
        
    
    Max Ramsay King's avatar
    Max Ramsay King committed
        def bin_to_subpol(self, subpol_bin):
            """
            Converts a binary string representation of a subpolicy to a subpolicy
    
            Parameters
            ------------
            subpol_bin -> str
                Binary representation of a subpolicy
    
            
            Returns
            -----------
            policy -> [(subpolicy)]
            """
    
    Max Ramsay King's avatar
    Max Ramsay King committed
            pol = []
            for idx in range(2):
    
    
    Max Ramsay King's avatar
    Max Ramsay King committed
                if subpol_bin[idx*12:(idx*12)+4] in self.bin_to_aug:
                    trans = self.bin_to_aug[subpol_bin[idx*12:(idx*12)+4]]
    
    Max Ramsay King's avatar
    Max Ramsay King committed
                else:
                    trans = random.choice(self.just_augs)
    
                mag_is_none = not self.aug_space_dict[trans]
    
    
    Max Ramsay King's avatar
    Max Ramsay King committed
                if subpol_bin[(idx*12)+4: (idx*12)+8] in self.bin_to_prob:
                    prob = float(self.bin_to_prob[subpol_bin[(idx*12)+4: (idx*12)+8]])
    
    Max Ramsay King's avatar
    Max Ramsay King committed
                else:
                    prob = float(random.randrange(0, 11, 1) / 10)
    
    
    Max Ramsay King's avatar
    Max Ramsay King committed
                if subpol_bin[(idx*12)+8:(idx*12)+12] in self.bin_to_mag:
                    mag = int(self.bin_to_mag[subpol_bin[(idx*12)+8:(idx*12)+12]])
    
    Max Ramsay King's avatar
    Max Ramsay King committed
                else:
                    mag = int(random.randrange(0, 10, 1))
    
                if mag_is_none:
                    mag = None
                pol.append((trans, prob, mag))
            pol = [tuple(pol)]
            return pol   
    
    
        def subpol_to_bin(self, subpol):
    
    Max Ramsay King's avatar
    Max Ramsay King committed
            """
            Converts a subpolicy to its binary representation 
    
            Parameters
            ------------
            subpol -> ((transforamtion, probability, magnitude), (trans., prob., mag.))
    
            Returns
            ------------
            bin_pol -> str
                Binary representation of the subpolicy
            """
            bin_pol = ''  
    
    Max Ramsay King's avatar
    Max Ramsay King committed
            trans1, prob1, mag1 = subpol[0]
            trans2, prob2, mag2 = subpol[1]
    
    
    Max Ramsay King's avatar
    Max Ramsay King committed
            bin_pol += self.aug_to_bin[trans1] + self.prob_to_bin[str(prob1)]
    
    Max Ramsay King's avatar
    Max Ramsay King committed
            
            if mag1 == None:
    
    Max Ramsay King's avatar
    Max Ramsay King committed
                bin_pol += '1111' 
    
    Max Ramsay King's avatar
    Max Ramsay King committed
            else:
    
    Max Ramsay King's avatar
    Max Ramsay King committed
                bin_pol += self.mag_to_bin[str(mag1)]
    
    Max Ramsay King's avatar
    Max Ramsay King committed
            
    
    Max Ramsay King's avatar
    Max Ramsay King committed
            bin_pol += self.aug_to_bin[trans2] + self.prob_to_bin[str(prob2)]
    
    Max Ramsay King's avatar
    Max Ramsay King committed
            
            if mag2 == None:
    
    Max Ramsay King's avatar
    Max Ramsay King committed
                bin_pol += '1111' 
    
    Max Ramsay King's avatar
    Max Ramsay King committed
            else:
    
    Max Ramsay King's avatar
    Max Ramsay King committed
                bin_pol += self.mag_to_bin[str(mag2)]
    
    Max Ramsay King's avatar
    Max Ramsay King committed
                
    
    Max Ramsay King's avatar
    Max Ramsay King committed
            return bin_pol
    
    Max Ramsay King's avatar
    Max Ramsay King committed
    
    
        def choose_parents(self, parents, parents_weights):
    
    Max Ramsay King's avatar
    Max Ramsay King committed
            """
            Chooses parents from which the next policy will be generated from
    
            Parameters
            ------------
            parents -> [policy, policy, ...]
    
            parents_weights -> [float, float, ...]
    
            Returns
            ------------
            (parent1, parent2) -> (policy, policy)
            
            """
    
    Max Ramsay King's avatar
    Max Ramsay King committed
            parent1 = random.choices(parents, parents_weights, k=1)[0][0]
            parent2 = random.choices(parents, parents_weights, k=1)[0][0]
            while parent2 == parent1:
                parent2 = random.choices(parents, parents_weights, k=1)[0][0]
            parent1 = self.subpol_to_bin(parent1)
            parent2 = self.subpol_to_bin(parent2)
            return (parent1, parent2)
    
        
        def generate_children(self):
    
    Max Ramsay King's avatar
    Max Ramsay King committed
            """
            Generates children via the random crossover method
    
            Returns 
            ------------
            new_pols -> [child_policy, child_policy, ...]
            """
    
    Max Ramsay King's avatar
    Max Ramsay King committed
            parent_acc = sorted(self.history, key = lambda x: x[1], reverse=True)
    
    Max Ramsay King's avatar
    Max Ramsay King committed
            parents = [x[0] for x in parent_acc]
            parents_weights = [x[1] for x in parent_acc]
            new_pols = []
            for _ in range(self.num_offspring):
                parent1, parent2 = self.choose_parents(parents, parents_weights)
    
    Max Ramsay King's avatar
    Max Ramsay King committed
                cross_over = random.randrange(1, int(len(parent2)/2), 1)
                cross_over2 = random.randrange(int(len(parent2)/2), int(len(parent2)), 1)
    
    Max Ramsay King's avatar
    Max Ramsay King committed
                child = parent1[:cross_over]
    
    Max Ramsay King's avatar
    Max Ramsay King committed
                child += parent2[cross_over:int(len(parent2)/2)]
                child += parent1[int(len(parent2)/2):int(len(parent2)/2)+cross_over2]
                child += parent2[int(len(parent2)/2)+cross_over2:]
    
    Max Ramsay King's avatar
    Max Ramsay King committed
                new_pols.append(child)
            return new_pols
    
        
    
    Max Ramsay King's avatar
    Max Ramsay King committed
        def learn(self, train_dataset, test_dataset, child_network_architecture, iterations = 100):
    
    Max Ramsay King's avatar
    Max Ramsay King committed
            """
            Generates policies through a genetic algorithm. 
    
            Parameters
            ------------
            train_dataset -> torchvision.dataset
    
            test_dataset -> torchvision.dataset
    
            child_network_architecture -> 
    
            iterations -> int
                number of iterations to run the instance for
            """
    
    Max Ramsay King's avatar
    Max Ramsay King committed
    
            for idx in range(iterations):
    
    Max Ramsay King's avatar
    Max Ramsay King committed
                print("ITERATION: ", idx)
                if len(self.history) < self.num_offspring:
    
    Max Ramsay King's avatar
    Max Ramsay King committed
                    policy = [self.gen_random_subpol()]
                else:
                    policy = self.bin_to_subpol(random.choice(self.generate_children()))
    
    Max Ramsay King's avatar
    Max Ramsay King committed
                print("Policy: ", policy)
    
    Max Ramsay King's avatar
    Max Ramsay King committed
                
                reward = self._test_autoaugment_policy(policy,
                                                    child_network_architecture,
                                                    train_dataset,
                                                    test_dataset)  
    
    Max Ramsay King's avatar
    Max Ramsay King committed
                print("reward: ", reward)