Python models.__dict__() Examples

def create_model(model_type,model_name,out_dim):
        # Create Similarity Prediction Network (SPN) by model surgery
        model = models.__dict__[model_type].__dict__[model_name](out_dim=out_dim)
        n_feat = model.last.in_features

        # Replace task-dependent module
        model.last = nn.Sequential(
            nn.Linear(n_feat*2, n_feat*4),
            nn.BatchNorm1d(n_feat*4),
            nn.ReLU(inplace=True),
            nn.Linear(n_feat*4, 2)
        )

        # Replace task-dependent function
        def new_logits(self, x):
            feat1, feat2 = PairEnum(x)
            featcat = torch.cat([feat1, feat2], 1)
            out = self.last(featcat)
            return out
        model.logits = MethodType(new_logits, model)

        return model 

def main():
    exp_dir = os.path.dirname(args.config)
    
    with open(args.config) as f:
        if version.parse(yaml.version >= "5.1"):
            config = yaml.load(f, Loader=yaml.FullLoader)
        else:
            config = yaml.load(f)

    for k, v in config.items():
        setattr(args, k, v)
    
    model = models.modules.__dict__[args.model['module']['arch']](args.model['module'])
    model = torch.nn.DataParallel(model)
    
    ckpt_path = exp_dir + '/checkpoints/ckpt_iter_{}.pth.tar'.format(args.iter)
    save_path = exp_dir + '/checkpoints/convert_iter_{}.pth.tar'.format(args.iter)
    ckpt = torch.load(ckpt_path)
    weight = ckpt['state_dict']
    model.load_state_dict(weight, strict=True)
    model = model.module.image_encoder
    
    torch.save(model.state_dict(), save_path) 

def main():
    args.distributed = args.world_size > 1
    args.gpu = 0
    if args.distributed:
        args.gpu = args.rank % torch.cuda.device_count()
        torch.cuda.set_device(args.gpu)
        dist.init_process_group(backend=args.dist_backend, init_method=args.dist_url, world_size=args.world_size)

    if args.fp16: assert torch.backends.cudnn.enabled, "fp16 mode requires cudnn backend to be enabled."

    model = models.__dict__[args.arch]().cuda()
    if args.distributed: model = DDP(model)

    data, train_sampler = torch_loader(f'{args.data}-sz/160', 128, 256)
    learner = Learner.from_model_data(model, data)
    learner.crit = F.cross_entropy
    learner.metrics = [accuracy, top5]
    if args.fp16: learner.half()
    wd=2e-5
    update_model_dir(learner, args.save_dir)
    fit(learner, '1', 0.03, 1, train_sampler, wd)

    data, train_sampler = torch_loader(f'{args.data}-sz/320', 128, 256)
    learner.set_data(data)
    fit(learner, '3', 1e-1, 1, train_sampler, wd)

    data, train_sampler = torch_loader(args.data, 128, 256)
    learner.set_data(data)
    fit(learner, '3', 1e-1, 1, train_sampler, wd)
    print('Finished!') 

def build_model():

    model = models.__dict__[args.arch](pretrained=True, num_classes=101)
    model.cuda()
    return model 

def build_model():

    model_name = args.modality + "_" + args.arch
    model = models.__dict__[model_name](pretrained=True, num_classes=101)
    if args.arch.startswith('alexnet') or args.arch.startswith('vgg'):
        model.features = torch.nn.DataParallel(model.features)
        model.cuda()
    else:
        model = torch.nn.DataParallel(model).cuda()
    return model 

def create_model(model_type,model_name,out_dim):
        # Prepare Constrained Clustering Network (CCN)
        model = models.__dict__[model_type].__dict__[model_name](out_dim=out_dim)
        return model 

def create_model(model_type,model_name,out_dim):
        # This function create the model for specific learner
        # The create_model(), forward_with_criterion(), and learn() are task-dependent
        # Do surgery to generic model if necessary
        model = models.__dict__[model_type].__dict__[model_name](out_dim=out_dim)
        #n_feat = model.last.in_features  # This information is useful
        return model 

def get_model(args):
    if args.first_layer_dense:
        args.first_layer_type = "DenseConv"

    print("=> Creating model '{}'".format(args.arch))
    model = models.__dict__[args.arch]()

    # applying sparsity to the network
    if (
        args.conv_type != "DenseConv"
        and args.conv_type != "SampleSubnetConv"
        and args.conv_type != "ContinuousSparseConv"
    ):
        if args.prune_rate < 0:
            raise ValueError("Need to set a positive prune rate")

        set_model_prune_rate(model, prune_rate=args.prune_rate)
        print(
            f"=> Rough estimate model params {sum(int(p.numel() * (1-args.prune_rate)) for n, p in model.named_parameters() if not n.endswith('scores'))}"
        )

    # freezing the weights if we are only doing subnet training
    if args.freeze_weights:
        freeze_model_weights(model)

    return model 

def main():
    parser = argparse.ArgumentParser(description='Compute Model FLOPs and PN')
    parser.add_argument('--arch', default='resnet50', type=str)
    parser.add_argument('--feature_dim', default=256, type=int)
    parser.add_argument('--input_size', default=112, type=int)
    parser.add_argument('--batch_size', default=1, type=int)
    args = parser.parse_args()

    model = models.__dict__[args.arch](feature_dim=args.feature_dim)
    flops, n_conv, n_linear = compute_flops(model, args.input_size,
                                            args.batch_size)
    bits = compute_param_number(model)

    print('[{} ({} conv, {} linear)] FLOPs: {:.2f} G, PN: {:.2f} M'.format(
        args.arch, n_conv, n_linear, flops, bits)) 

def __init__(self, model_path, device=None, benchmark: bool=True):
        '''
        Args:
            module_path: path to pre-trained model (available to download, see README)
            device: CUDA device to use. str or torch.device instance
                warning: this is restricted to 'cpu' or 'cuda' only
                    ('cuda:1' won't work due to main package arcitecture)
                default is choose 'cuda' if available
            benchmark: to enable cudnn benchmark mode or not
        '''
        self.model_path = model_path
        if device is None:
            device = 'cuda' if torch.cuda.is_available() else 'cpu'
        self.device = device

        if benchmark:
            torch.backends.cudnn.enabled = True
            torch.backends.cudnn.benchmark = True

        snapshot = torch.load(self.model_path, map_location=self.device)
        self.param = snapshot['args']

        self.transform  = transforms.Compose([
            transforms.PreCrop(self.param.pre_crop_expand),
            transforms.TrainScale2WH((self.param.crop_width, self.param.crop_height)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
        ])

        self.net = models.__dict__[self.param.arch](self.param.modelconfig, None)
        self.net.train(False).to(self.device)

        weights = models.remove_module_dict(snapshot['state_dict'])
        self.net.load_state_dict(weights) 

def init_optimizer(self):
        optimizer_arg = {'params':self.model.parameters(),
                         'lr':self.config['lr'],
                         'weight_decay':self.config['weight_decay']}
        if self.config['optimizer'] in ['SGD','RMSprop']:
            optimizer_arg['momentum'] = self.config['momentum']
        elif self.config['optimizer'] in ['Rprop']:
            optimizer_arg.pop('weight_decay')
        elif self.config['optimizer'] == 'amsgrad':
            optimizer_arg['amsgrad'] = True
            self.config['optimizer'] = 'Adam'

        self.optimizer = torch.optim.__dict__[self.config['optimizer']](**optimizer_arg)
        self.scheduler = torch.optim.lr_scheduler.MultiStepLR(self.optimizer, milestones=self.config['schedule'],
                                                              gamma=0.1) 

def load_snapshot(snapshot_file):
    """Load a training snapshot"""
    print("--- Loading model from snapshot")

    # Create network
    norm_act = partial(InPlaceABN, activation="leaky_relu", activation_param=.01)
    body = models.__dict__["net_wider_resnet38_a2"](norm_act=norm_act, dilation=(1, 2, 4, 4))
    head = DeeplabV3(4096, 256, 256, norm_act=norm_act, pooling_size=(84, 84))

    # Load snapshot and recover network state
    data = torch.load(snapshot_file)
    body.load_state_dict(data["state_dict"]["body"])
    head.load_state_dict(data["state_dict"]["head"])

    return body, head, data["state_dict"]["cls"] 

def _optimizer_options(parser):
    """
    Options specific for optimizers
    """
    # List of possible optimizers already implemented in PyTorch
    optimizer_options = [name for name in torch.optim.__dict__ if callable(torch.optim.__dict__[name])]

    parser_optimizer = parser.add_argument_group('OPTIMIZER', 'Optimizer Options')

    parser_optimizer.add_argument('--optimizer-name',
                                  choices=optimizer_options,
                                  default='SGD',
                                  help='optimizer to be used for training')
    parser_optimizer.add_argument('--lr',
                                  type=float,
                                  default=0.001,
                                  help='learning rate to be used for training')
    parser_optimizer.add_argument('--decay-lr',
                                  type=int,
                                  default=None,
                                  help='drop LR by 10 every N epochs')
    parser_optimizer.add_argument('--momentum',
                                  type=float,
                                  default=0,
                                  help='momentum (parameter for the optimizer)')
    parser_optimizer.add_argument('--dampening',
                                  type=float,
                                  default=0,
                                  help='dampening (parameter for the SGD)')
    parser_optimizer.add_argument('--weight-decay',
                                  type=float,
                                  default=0,
                                  help='weight_decay coefficient, also known as L2 regularization') 

def create_model(self):
        cfg = self.config

        # Define the backbone (MLP, LeNet, VGG, ResNet ... etc) of model
        model = models.__dict__[cfg['model_type']].__dict__[cfg['model_name']]()

        # Apply network surgery to the backbone
        # Create the heads for tasks (It can be single task or multi-task)
        n_feat = model.last.in_features

        # The output of the model will be a dict: {task_name1:output1, task_name2:output2 ...}
        # For a single-headed model the output will be {'All':output}
        model.last = nn.ModuleDict()
        for task,out_dim in cfg['out_dim'].items():
            model.last[task] = nn.Linear(n_feat,out_dim)

        # Redefine the task-dependent function
        def new_logits(self, x):
            outputs = {}
            for task, func in self.last.items():
                outputs[task] = func(x)
            return outputs

        # Replace the task-dependent function
        model.logits = MethodType(new_logits, model)
        # Load pre-trained weights
        if cfg['model_weights'] is not None:
            print('=> Load model weights:', cfg['model_weights'])
            model_state = torch.load(cfg['model_weights'],
                                     map_location=lambda storage, loc: storage)  # Load to CPU.
            model.load_state_dict(model_state)
            print('=> Load Done')
        return model 

def _get_criterion(criterion_name, disable_databalancing, dataset_folder, inmem, workers, **kwargs):
    """
    This function serves as an interface between the command line and the criterion.

    Parameters
    ----------
     criterion_name : string
        Name of the criterion
    disable_databalancing : boolean
        If True the criterion will not be fed with the class frequencies. Use with care.
    dataset_folder : String
        Location of the dataset on the file system
    inmem : boolean
        Load the whole dataset in memory. If False, only file names are stored and images are loaded
        on demand. This is slower than storing everything in memory.
    workers : int
        Number of workers to use for the dataloaders

    Returns
    -------
    torch.nn
        The initalized criterion

    """
    # Verify that the criterion exists
    assert criterion_name in torch.nn.__dict__

    # Instantiate the criterion
    criterion = torch.nn.__dict__[criterion_name]()

    if not disable_databalancing:
        try:
            logging.info('Loading weights for data balancing')
            weights = _load_class_frequencies_weights_from_file(dataset_folder, inmem, workers, **kwargs)
            criterion.weight = torch.from_numpy(weights).type(torch.FloatTensor)
        except:
            logging.warning('Unable to load information for data balancing. Using normal criterion')
    return criterion 

def _get_optimizer(optimizer_name, model, **kwargs):
    """
    This function serves as interface between the command line and the optimizer.
    In fact each optimizer has a different set of parameters and in this way one can just change the optimizer
    in his experiments just by changing the parameters passed to the entry point.

    Parameters
    ----------
    optimizer_name:
        Name of the optimizers. See: torch.optim for a list of possible values
    model:
        The model with which the training will be done
    kwargs:
        List of all arguments to be used to init the optimizer
    Returns
    -------
    torch.optim
        The optimizer initialized with the provided parameters
    """
    # Verify the optimizer exists
    assert optimizer_name in torch.optim.__dict__

    params = {}
    # For all arguments declared in the constructor signature of the selected optimizer
    for p in inspect.getfullargspec(torch.optim.__dict__[optimizer_name].__init__).args:
        # Add it to a dictionary in case it exists a corresponding value in kwargs
        if p in kwargs:
            params.update({p: kwargs[p]})
    # Create an return the optimizer with the correct list of parameters
    return torch.optim.__dict__[optimizer_name](model.parameters(), **params) 

def main():
    global args
    args = parser.parse_args()
    model = models.__dict__[args.arch]()    
    print(model)
    input = torch.randn(1, 3, 224, 224)
    model.train()
    # model.eval()
    flops, params = profile(model, inputs=(input, ))
    print("flops = ", flops)
    print("params = ", params)
    flops, params = clever_format([flops, params], "%.3f")
    print("flops = ", flops)
    print("params = ", params) 

def main():
    args.distributed = args.world_size > 1
    args.gpu = 0
    if args.distributed: args.gpu = args.rank % torch.cuda.device_count()

    if args.distributed:
        torch.cuda.set_device(args.gpu)
        dist.init_process_group(backend=args.dist_backend, init_method=args.dist_url,
                                world_size=args.world_size)

    if args.fp16: assert torch.backends.cudnn.enabled, "missing cudnn"

    model = cifar10models.__dict__[args.arch] if args.arch in cifar10_names else models.__dict__[args.arch]
    if args.pretrained: model = model(pretrained=True)
    else: model = model()

    model = model.cuda()
    if args.distributed: model = DDP(model)
    if args.data_parallel: model = nn.DataParallel(model, [0,1,2,3])

    data, train_sampler = torch_loader(args.data, args.sz)

    learner = Learner.from_model_data(model, data)
    #print (learner.summary()); exit()
    learner.crit = F.cross_entropy
    learner.metrics = [accuracy]
    if args.fp16: learner.half()

    if args.prof: args.epochs,args.cycle_len = 1,0.01
    if args.use_clr: args.use_clr = tuple(map(float, args.use_clr.split(',')))

    # Full size
    update_model_dir(learner, args.save_dir)
    sargs = save_args('first_run', args.save_dir)

    if args.warmup:
        learner.fit(args.lr/10, 1, cycle_len=1, sampler=train_sampler, wds=args.weight_decay,
                use_clr_beta=(100,1,0.9,0.8), loss_scale=args.loss_scale, **sargs)

    learner.fit(args.lr,args.epochs, cycle_len=args.cycle_len,
                sampler=train_sampler, wds=args.weight_decay,
                use_clr_beta=args.use_clr, loss_scale=args.loss_scale,
                **sargs)
    save_sched(learner.sched, args.save_dir)

    print('Finished!')

    if args.use_tta:
        log_preds,y = learner.TTA()
        preds = np.mean(np.exp(log_preds),0)
        acc = accuracy(torch.FloatTensor(preds),torch.LongTensor(y))
        print('TTA acc:', acc)

        with open(f'{args.save_dir}/tta_accuracy.txt', "a", 1) as f:
            f.write(time.strftime("%Y-%m-%dT%H:%M:%S")+f"\tTTA accuracty: {acc}\n") 

def main():
    args, cuda = parse_arg()
    logfile, run = config_envs(args, cuda)
    train_loader, val_loader = prepare_data(args, cuda)

    net = models.__dict__[args.model]()
    criterion = torch.nn.CrossEntropyLoss()

    if cuda:
        ctx = [int(id) for id in args.gpus.split(',')]
        net = torch.nn.DataParallel(net, device_ids=ctx)
        cudnn.benchmark = True
        net, criterion = net.cuda(), criterion.cuda()
    # early stopping parameters
    patience = args.patience
    best_loss = 1e4

    # Print model to logfile
    print(net, file=logfile)

    # Change optimizer for finetuning
    optimizer = optim.Adam(net.parameters())

    for e in range(args.nepochs):
        start = time.time()
        train_loss, train_acc = train(net, train_loader,
            criterion, optimizer, cuda)
        val_loss, val_acc = validate(net, val_loader, criterion, cuda)
        end = time.time()

        # print stats
        stats ="""Epoch: {}\t train loss: {:.3f}, train acc: {:.3f}\t
                val loss: {:.3f}, val acc: {:.3f}\t
                time: {:.1f}s""".format( e, train_loss, train_acc, val_loss,
                val_acc, end-start)
        print(stats)
        print(stats, file=logfile)
        # tensorboard feature
        # log_value('train_loss', train_loss, e)
        # log_value('val_loss', val_loss, e)
        # log_value('train_acc', train_acc, e)
        # log_value('val_acc', val_acc, e)

        #early stopping and save best model
        if val_loss < best_loss:
            best_loss = val_loss
            patience = args.patience
            utils.save_model({
                'arch': args.model,
                'state_dict': net.state_dict()
            }, 'saved-models/{}-run-{}.pth.tar'.format(args.model, run))
        else:
            patience -= 1
            if patience == 0:
                print('Early stopped at epoch ' + str(e))
                break 

def main():
    global args
    args = parser.parse_args()

    assert args.output_path.endswith('.npy')

    print("=> creating model '{}'".format(args.arch))
    model = models.__dict__[args.arch](feature_dim=args.feature_dim)
    model = IdentityMapping(model)

    model.cuda()
    model = torch.nn.DataParallel(model).cuda()

    if args.load_path:
        if args.strict:
            classifier_keys = ['module.logits.weight', 'module.logits.bias']
            load_ckpt(args.load_path,
                      model,
                      ignores=classifier_keys,
                      strict=True)
        else:
            load_ckpt(args.load_path, model, strict=False)

    cudnn.benchmark = True

    normalize = transforms.Normalize(mean=[0.5, 0.5, 0.5],
                                     std=[0.25, 0.25, 0.25])

    test_loader = DataLoader(BinDataset(
        args.bin_file,
        transforms.Compose([
            transforms.Resize(args.input_size),
            transforms.ToTensor(),
            normalize,
        ])),
                             batch_size=args.batch_size,
                             shuffle=False,
                             num_workers=args.workers,
                             pin_memory=True)

    features = extract(test_loader, model)
    assert features.shape[1] == args.feature_dim

    print('saving extracted features to {}'.format(args.output_path))
    folder = os.path.dirname(args.output_path)
    if folder != '' and not os.path.exists(folder):
        os.makedirs(folder)
    np.save(args.output_path, features) 

def main():
    global args
    args = parser.parse_args()
    # Setup device
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    # Setup Model
    from collections import namedtuple
    Config = namedtuple('Config', ['iSz', 'oSz', 'gSz', 'batch'])
    config = Config(iSz=160, oSz=56, gSz=112, batch=1)  # default for training

    model = (models.__dict__[args.arch](config))
    model = load_pretrain(model, args.resume)
    model = model.eval().to(device)

    scales = [2**i for i in range_end(args.si, args.sf, args.ss)]
    meanstd = {'mean':[0.485, 0.456, 0.406], 'std':[0.229, 0.224, 0.225]}
    infer = Infer(nps=args.nps, scales=scales, meanstd=meanstd, model=model, device=device)

    print('| start'); tic = time.time()
    im = np.array(Image.open(args.img).convert('RGB'), dtype=np.float32)
    h, w = im.shape[:2]
    img = np.expand_dims(np.transpose(im, (2, 0, 1)), axis=0).astype(np.float32)
    img = torch.from_numpy(img / 255.).to(device)
    infer.forward(img)
    masks, scores = infer.getTopProps(.2, h, w)
    toc = time.time() - tic
    print('| done in %05.3f s' % toc)

    for i in range(masks.shape[2]):
        res = im[:,:,::-1].copy().astype(np.uint8)
        res[:, :, 2] = masks[:, :, i] * 255 + (1 - masks[:, :, i]) * res[:, :, 2]

        mask = masks[:, :, i].astype(np.uint8)
        _, contours, _ = cv2.findContours(mask.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
        cnt_area = [cv2.contourArea(cnt) for cnt in contours]
        cnt_max_id = np.argmax(cnt_area)
        contour = contours[cnt_max_id]
        polygons = contour.reshape(-1, 2)

        predict_box = cv2.boundingRect(polygons)
        predict_rbox = cv2.minAreaRect(polygons)
        rbox = cv2.boxPoints(predict_rbox)
        print('Segment Proposal Score: {:.3f}'.format(scores[i]))

        res = cv2.rectangle(res, (predict_box[0], predict_box[1]),
                      (predict_box[0]+predict_box[2], predict_box[1]+predict_box[3]), (0, 255, 0), 3)
        res = cv2.polylines(res, [np.int0(rbox)], True, (0, 255, 255), 3)
        cv2.imshow('Proposal', res)
        cv2.waitKey(0) 

def main():
    global args, best_acc, tb_writer, logger
    args = parser.parse_args()

    init_log('global', logging.INFO)

    if args.log != "":
        add_file_handler('global', args.log, logging.INFO)

    logger = logging.getLogger('global')
    logger.info(args)

    cfg = load_config(args)

    logger.info("config \n{}".format(json.dumps(cfg, indent=4)))
    
    logger.info("\n" + collect_env_info())

    if args.log_dir:
        tb_writer = SummaryWriter(args.log_dir)
    else:
        tb_writer = Dummy()

    # build dataset
    train_loader, val_loader = build_data_loader(cfg)

    if args.arch == 'Custom':
        from custom import Custom
        model = Custom(pretrain=True, anchors=cfg['anchors'])
    else:
        model = models.__dict__[args.arch](anchors=cfg['anchors'])

    logger.info(model)

    if args.pretrained:
        model = load_pretrain(model, args.pretrained)

    model = model.cuda()
    dist_model = torch.nn.DataParallel(model, list(range(torch.cuda.device_count()))).cuda()

    if args.resume and args.start_epoch != 0:
        model.features.unfix((args.start_epoch - 1) / args.epochs)

    optimizer, lr_scheduler = build_opt_lr(model, cfg, args, args.start_epoch)
    logger.info(lr_scheduler)
    # optionally resume from a checkpoint
    if args.resume:
        assert os.path.isfile(args.resume), '{} is not a valid file'.format(args.resume)
        model, optimizer, args.start_epoch, best_acc, arch = restore_from(model, optimizer, args.resume)
        dist_model = torch.nn.DataParallel(model, list(range(torch.cuda.device_count()))).cuda()
        epoch = args.start_epoch
        if dist_model.module.features.unfix(epoch/args.epochs):
            logger.info('unfix part model.')
            optimizer, lr_scheduler = build_opt_lr(dist_model.module, cfg, args, epoch)
        lr_scheduler.step(epoch)
        cur_lr = lr_scheduler.get_cur_lr()
        logger.info('epoch:{} resume lr {}'.format(epoch, cur_lr))

    logger.info('model prepare done')

    train(train_loader, dist_model, optimizer, lr_scheduler, args.start_epoch, cfg) 

def main():
    args.distributed = args.world_size > 1
    args.gpu = 0
    if args.distributed: args.gpu = args.rank % torch.cuda.device_count()

    if args.distributed:
        torch.cuda.set_device(args.gpu)
        dist.init_process_group(backend=args.dist_backend, init_method=args.dist_url,
                                world_size=args.world_size)

    if args.fp16: assert torch.backends.cudnn.enabled, "missing cudnn"

    model = cifar10models.__dict__[args.arch] if args.arch in cifar10_names else models.__dict__[args.arch]
    if args.pretrained: model = model(pretrained=True)
    else: model = model()

    model = model.cuda()
    if args.distributed: model = DDP(model)
    if args.data_parallel: model = nn.DataParallel(model, [0,1,2,3])

    data = torch_loader(args.data, args.sz)

    learner = Learner.from_model_data(model, data)
    #print (learner.summary()); exit()
    learner.crit = F.cross_entropy
    learner.metrics = [accuracy]
    if args.fp16: learner.half()

    if args.prof: args.epochs,args.cycle_len = 1,0.01
    if args.use_clr: args.use_clr = tuple(map(float, args.use_clr.split(',')))

    # Full size
    update_model_dir(learner, args.save_dir)
    sargs = save_args('first_run', args.save_dir)

    if args.warmup:
        learner.fit(args.lr/10, 1, cycle_len=1, wds=args.weight_decay,
                use_clr_beta=(100,1,0.9,0.8), loss_scale=args.loss_scale, **sargs)

    learner.fit(args.lr,args.epochs, cycle_len=args.cycle_len, wds=args.weight_decay,
                use_clr_beta=args.use_clr, loss_scale=args.loss_scale,
                **sargs)
    save_sched(learner.sched, args.save_dir)

    print('Finished!')

    if args.use_tta:
        log_preds,y = learner.TTA()
        preds = np.mean(np.exp(log_preds),0)
        acc = accuracy(torch.FloatTensor(preds),torch.LongTensor(y))
        print('TTA acc:', acc)

        with open(f'{args.save_dir}/tta_accuracy.txt', "a", 1) as f:
            f.write(time.strftime("%Y-%m-%dT%H:%M:%S")+f"\tTTA accuracty: {acc}\n") 

def main():
    args.distributed = args.world_size > 1
    args.gpu = 0
    if args.distributed: args.gpu = args.rank % torch.cuda.device_count()

    if args.distributed:
        torch.cuda.set_device(args.gpu)
        dist.init_process_group(backend=args.dist_backend, init_method=args.dist_url, world_size=args.world_size)

    if args.fp16: assert torch.backends.cudnn.enabled, "fp16 mode requires cudnn backend to be enabled."

    # create model
    if args.pretrained: model = models.__dict__[args.arch](pretrained=True)
    else:               model = models.__dict__[args.arch]()

    model = model.cuda()
    if args.distributed: model = DDP(model)

    data, train_sampler = torch_loader(f'{args.data}-160', 128, 256)
    learner = Learner.from_model_data(model, data)
    learner.crit = F.cross_entropy
    learner.metrics = [accuracy, top5]
    if args.fp16: learner.half()
    update_model_dir(learner, args.save_dir)

    wd=1e-4
    lr=0.1
    data, train_sampler = torch_loader(args.data, 224, 192)
    learner.set_data(data)
    fit(learner, '1', lr/4, 1, train_sampler, wd)
    fit(learner, '1', lr/2, 1, train_sampler, wd)
    fit(learner, '2', lr, 28, train_sampler, wd)

    #data, train_sampler = torch_loader(args.data, 224, 192)
    #learner.set_data(data)
    #fit(learner, '3', lr, 5, train_sampler, wd)
    fit(learner, '4', lr/10, 25, train_sampler, wd)
    fit(learner, '5', lr/100, 25, train_sampler, wd)

    data, train_sampler = torch_loader(args.data, 288, 128, min_scale=0.5)
    learner.set_data(data)
    fit(learner, '6', lr/500, 10, train_sampler, wd)
    #save_sched(learner.sched, args.save_dir)
    #fit(learner, '7', 1e-4, 10, train_sampler, wd/4)

    # TTA works ~50% of the time. Hoping top5 works better
    print('\n TTA \n')
    log_preds,y = learner.TTA()
    preds = np.mean(np.exp(log_preds),0)
    acc = accuracy(torch.FloatTensor(preds),torch.LongTensor(y))
    t5 = top5(torch.FloatTensor(preds),torch.LongTensor(y))
    print('TTA acc:', acc)
    print('TTA top5:', t5[0])
    with open(f'{args.save_dir}/tta_accuracy.txt', "a", 1) as f:
        f.write(time.strftime("%Y-%m-%dT%H:%M:%S")+f"\tTTA accuracy: {acc}\tTop5: {t5}")

    print('Finished!') 

def main():
    global args, conf
    args = parser.parse_args()

    torch.cuda.set_device(args.local_rank)

    try:
        world_size = int(os.environ['WORLD_SIZE'])
        distributed = world_size > 1
    except:
        distributed = False
        world_size = 1

    if distributed:
        dist.init_process_group(backend=args.dist_backend, init_method='env://')

    # Load configuration
    conf = config.load_config(args.config)

    # Create model
    model_params = utils.get_model_params(conf["network"])
    model = models.__dict__["net_" + conf["network"]["arch"]](**model_params)
    model.cuda()
    if distributed:
        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank],
                                                          output_device=args.local_rank)
    else:
        model = SingleGPU(model)

    # Resume from checkpoint
    checkpoint = torch.load(args.checkpoint)
    model.load_state_dict(checkpoint['state_dict'])

    # Data loading code
    valdir = os.path.join(args.data, 'val')
    val_transforms = utils.create_test_transforms(conf["input"], args.crop, args.scale, args.ten_crops)

    batch_size = conf["optimizer"]["batch_size"] if not args.ten_crops else conf["optimizer"]["batch_size"] // 10
    dataset = datasets.ImageFolder(valdir, transforms.Compose(val_transforms))
    val_loader = torch.utils.data.DataLoader(
        dataset, batch_size=batch_size // world_size, shuffle=False, sampler=TestDistributedSampler(dataset),
        num_workers=args.workers, pin_memory=True)

    criterion = nn.CrossEntropyLoss().cuda()
    utils.validate(val_loader, model, criterion, args.ten_crops, args.print_freq) 

def _training_options(parser):
    """
    Training options
    """
    # List of possible custom models already implemented
    # NOTE: If a model is missing and you get a argument parser error: check in the init file of models if its there!
    model_options = [name for name in models.__dict__ if callable(models.__dict__[name])]

    parser_train = parser.add_argument_group('TRAIN', 'Training Options')
    parser_train.add_argument('--model-name',
                              type=str,
                              choices=model_options,
                              default='CNN_basic',
                              help='which model to use for training')
    parser_train.add_argument('--batch-size',
                              type=int,
                              default=64,
                              help='input batch size for training')
    parser_train.add_argument('--epochs',
                              type=int,
                              default=5,
                              help='how many epochs to train')
    parser_train.add_argument('--pretrained',
                              action='store_true',
                              default=False,
                              help='use pretrained model. (Not applicable for all models)')
    parser_train.add_argument('--load-model',
                              type=str,
                              default=None,
                              help='path to latest checkpoint')
    parser_train.add_argument('--resume',
                              type=str,
                              default=None,
                              help='path to latest checkpoint')
    parser_train.add_argument('--start-epoch',
                              type=int,
                              metavar='N',
                              default=0,
                              help='manual epoch number (useful on restarts)')
    parser_train.add_argument('--validation-interval',
                              type=int,
                              default=1,
                              help='run evaluation on validation set every N epochs')
    parser_train.add_argument('--checkpoint-all-epochs',
                              action='store_true',
                              default=False,
                              help='make a checkpoint after every epoch') 

def prepare(cls, model_name, **kwargs):
        """
        Loads and prepares the data, the optimizer and the criterion

        Parameters
        ----------
        model_name : str
            Name of the model. Used for loading the model.
        kwargs : dict
            Any additional arguments.

        Returns
        -------
        model : DataParallel
            The model to train
        num_classes : int
            How many different classes there are in our problem. Used for loading the model.
        best_value : float
            Best value of the model so far. Non-zero only in case of --resume being used
        train_loader : torch.utils.data.dataloader.DataLoader
            Training dataloader
        val_loader : torch.utils.data.dataloader.DataLoader
            Validation dataloader
        test_loader : torch.utils.data.dataloader.DataLoader
            Test set dataloader
        optimizer : torch.optim
            Optimizer to use during training, e.g. SGD
        criterion : torch.nn.modules.loss
            Loss function to use, e.g. cross-entropy
        """
        # Get the selected model input size
        model_expected_input_size = models.__dict__[model_name]().expected_input_size
        if type(model_expected_input_size) is not tuple or len(model_expected_input_size) != 2:
            logging.error('Model {model_name} expected input size is not a tuple. '
                          'Received: {model_expected_input_size}'
                          .format(model_name=model_name,
                                  model_expected_input_size=model_expected_input_size))
            sys.exit(-1)
        logging.info('Model {} expects input size of {}'.format(model_name, model_expected_input_size))

        # Setting up the dataloaders
        train_loader, val_loader, test_loader, num_classes = set_up_dataloaders(model_expected_input_size, **kwargs)

        # Setting up model, optimizer, criterion
        model, criterion, optimizer, best_value = set_up_model(model_name=model_name, num_classes=num_classes,
                                                               **kwargs)
        return  model, num_classes, best_value, train_loader, val_loader, test_loader, optimizer, criterion