示例#1
0
class Trainer:
    def __init__(self):
        self._setup_logger()
        torch.backends.cudnn.deterministic = True
        self.device = torch.device(
            "cuda:%d" % (args.gpu) if torch.cuda.is_available() else "cpu")
        torch.cuda.set_device(args.gpu)
        self.dataset = Cifar10()
        self.saver = Saver(self.save_dir)
        if args.arch == 'auto':
            self.Net = lambda: AutoNetwork(args.init_channels, self.dataset.
                                           num_classes, args.layers,
                                           nn.CrossEntropyLoss())
        elif args.arch == 'VGG11':
            self.Net = lambda: VGG('VGG11', self.dataset.num_classes,
                                   nn.CrossEntropyLoss())
        elif args.arch == 'mobilenet':
            self.Net = lambda: MobileNetV2(self.dataset.num_classes,
                                           nn.CrossEntropyLoss())
        elif args.arch == 'simple':
            self.Net = lambda: SimpleNetwork(self.dataset.num_classes,
                                             nn.CrossEntropyLoss())
        else:
            raise Exception('Net not defined!')

    def _setup_logger(self):
        self.save_dir = os.path.join('./checkpoints', args.logdir)
        if not os.path.isdir(self.save_dir):
            os.makedirs(self.save_dir)
        log_path = os.path.join(
            self.save_dir,
            datetime.datetime.now().strftime('%Y-%m-%d-%H-%M-%S'))
        self.logger = L.create_logger(args.logdir, log_path)
        for arg in vars(args):
            self.logger.info("%-25s %-20s" % (arg, getattr(args, arg)))

    def test_beta(self, optimizee):
        self.logger.info('Testing beta')
        torch.manual_seed(2 * args.seed)
        optimizee.reset_model_parameters()
        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
            optimizee.optimizer,
            float(args.max_epoch),
            eta_min=args.learning_rate_min)
        optimizee.optimizer.print_beta_greedy()
        for epoch in range(0, args.max_test_epoch):
            scheduler.step()
            lr = scheduler.get_lr()[0]
            self.train_epoch(epoch, lr, optimizee, False)
            if (epoch + 1) % args.test_freq == 0:
                acc = self.eval(optimizee.model)

    def train(self):
        optimizee = Optimizee(self.Net)
        start_episode, start_epoch, best_test_acc = 0, 0, 0.
        if args.checkpoint != '':
            start_episode, start_epoch = self.saver.load_checkpoint(
                optimizee, args.checkpoint)
        # Hao: in my implementation, the optimizers for alpha and beta will preserve their states across episodes
        for episode in range(start_episode, args.max_episodes):
            torch.manual_seed(args.seed + episode)
            optimizee.reset_model_parameters()
            if args.arch == 'auto':
                optimizee.reset_arch_parameters()
            scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
                optimizee.optimizer,
                float(args.max_epoch),
                eta_min=args.learning_rate_min,
                last_epoch=start_epoch - 1)
            for epoch in range(
                    start_epoch,
                    args.max_epoch):  # loop over the dataset multiple times
                scheduler.step()
                lr = scheduler.get_lr()[0]
                training_status = self.train_epoch(epoch, lr, optimizee)
                if (epoch + 1) % args.test_freq == 0:
                    acc = self.eval(optimizee.model)
                    # if accuracy not above random, discard this example and start a new one
                    if acc <= 0.11 or training_status == False:
                        self.logger.warning(
                            'training_status false or acc too low, break')
                        break
                    checkpoint_path = self.save_dir + '/episode_{}_epoch_{}_acc_{}'.format(
                        str(episode), str(epoch), str(acc))
                    self.saver.save_checkpoint(optimizee, epoch, episode,
                                               checkpoint_path)
        return optimizee

    def train_epoch(self, epoch, lr, optimizee, train_beta=True):
        if train_beta:
            optimizee.optimizer.beta_temp = np.interp(
                epoch + 1, [0, args.max_epoch],
                [args.min_beta_temp, args.max_beta_temp])
            status = self._train_epoch(epoch, lr, optimizee)
            self.saver.write_beta_embedding()
            optimizee.optimizer.print_beta_greedy()
        else:
            optimizee.optimizer.beta_temp = args.max_beta_temp
            status = self._train_epoch_fix_beta(epoch, lr, optimizee)
        return status

    def eval(self, model):
        correct, total = 0, 0
        with torch.no_grad():
            for images, labels in self.dataset.test_loader:
                images, labels = images.to(self.device), labels.to(self.device)
                outputs = model(images)
                _, predicted = torch.max(outputs.data, 1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()
        self.logger.info(
            'Accuracy of the network on the 10000 test images: %d %%' %
            (100 * correct / total))
        return float(correct) / total

    def _train_epoch(self, epoch, lr, optimizee):
        losses, next_step_losses, meta_update_losses = 0.0, 0.0, 0.0
        model_grad_norms, beta_grad_norms = 0., 0.
        for i, data in enumerate(self.dataset.train_loader, 0):
            x_train, y_train = data[0].to(self.device), data[1].to(self.device)
            val_data = next(iter(self.dataset.val_loader))
            x_val, y_val = val_data[0].to(self.device), val_data[1].to(
                self.device)

            # Derive \frac{\partial L}{\partial \theta}
            optimizee.model.zero_grad()
            loss = optimizee.model.forward_pass(x_train, y_train)
            loss.backward()
            model_grad_norms += nn.utils.clip_grad_norm_(
                optimizee.model.parameters(), args.model_grad_norm)
            param_updates = optimizee.optimizer.step()

            # do a differentiable step of update over optimizee.symbolic_model
            optimizee.differentiable_update(param_updates)

            # check the next-step loss after update using the symbolic model
            # so that it is on the computational graph
            # TODO use validation data? not sure
            next_step_loss = optimizee.symbolic_model.forward_pass(
                x_train, y_train)
            if (math.isnan(next_step_loss.item())):
                self.logger.error('next step loss becomes NaN, break')
                raise Exception
            next_step_losses += next_step_loss.item()
            losses += loss.item()
            meta_update_losses += next_step_loss

            # do a non-differentiable update over optimizee.model if next_step_loss is smaller
            if (next_step_loss.item() <
                    loss.item()):  # This is still important for performance
                optimizee.update(param_updates)

            # forsee bptt_steps then update beta
            if (i + 1) % args.bptt_step == 0:
                beta_grad_norms += optimizee.beta_step(
                    meta_update_losses).item()
                # let saver save the grads for beta
                with torch.no_grad():
                    self.saver.add_beta_grads([
                        b.grad.abs().mean() for b in optimizee.optimizer.beta
                    ])
                meta_update_losses = 0
                optimizee.sync_symbolic_model()

            if args.arch_training and (i + 1) % args.update_alpha_step == 0:
                optimizee.alpha_step(
                    x_train, y_train, x_val, y_val,
                    lr)  # TODO: make sure alpha step won't change weights
                optimizee.sync_symbolic_model(skip_weights=True)

            if (i + 1) % args.log_freq == 0:  # print every 2000 mini-batches
                beta_out = optimizee.optimizer.beta[-1].data.cpu().numpy()
                alpha_out = optimizee.model.alphas_normal[-1].data.cpu().numpy(
                ) if args.arch_training else ''
                self.logger.info(
                    '[%d, %5d] loss: %.3f/%.3f, next step loss: %.3f, beta[-1]/L2(g): %s/%s, alpha[-1]: %s'
                    % (epoch + 1, i + 1, losses / args.log_freq,
                       model_grad_norms / args.log_freq,
                       next_step_losses / args.log_freq, beta_out,
                       beta_grad_norms / args.log_freq, alpha_out))
                losses, next_step_losses, model_grad_norms, beta_grad_norms = 0., 0., 0., 0.
                self.saver.write_beta(optimizee.optimizer.beta)
        return True

    def _train_epoch_fix_beta(self, epoch, lr, optimizee):
        losses, model_grad_norms = 0.0, 0.0
        for i, data in enumerate(self.dataset.train_loader, 0):
            x_train, y_train = data[0].to(self.device), data[1].to(self.device)
            val_data = next(iter(self.dataset.val_loader))
            x_val, y_val = val_data[0].to(self.device), val_data[1].to(
                self.device)
            optimizee.model.zero_grad()
            loss = optimizee.model.forward_pass(x_train, y_train)
            loss.backward()
            model_grad_norms += nn.utils.clip_grad_norm_(
                optimizee.model.parameters(), args.model_grad_norm)
            param_updates = optimizee.optimizer.step(do_update=True)
            losses += loss.item()
            if args.use_darts_arch and args.train_alpha:
                optimizee.alpha_step(x_train, y_train, x_val, y_val, lr)
            if (i + 1) % args.log_freq == 0:  # print every 2000 mini-batches
                beta_out = optimizee.optimizer.beta[-1].data.cpu().numpy()
                alpha_out = optimizee.model.alphas_normal[-1].data.cpu().numpy(
                ) if args.arch_training else 'None'
                self.logger.info(
                    '[%d, %5d] loss: %.3f/%.3f, beta[-1]: %s, alpha[-1]: %s' %
                    (epoch + 1, i + 1, losses / args.log_freq,
                     model_grad_norms / args.log_freq, beta_out, alpha_out))
                losses, model_grad_norms = 0., 0.
        return True