示例#1
0
def main():
    start_time = time.time()

    in_arg = args_input()
    # get the data
    dataset, dataloaders = functions.load_data()

    # get the model
    model = functions.get_model(in_arg.arch)

    # get the classifier, criterion, optimizer, device
    model, classifier, criterion, optimizer, device = network_param(
        model, in_arg.arch, in_arg.hidden_units, in_arg.learning_rate,
        in_arg.gpu)

    model.to(device)

    print('Training model...')
    functions.train_model(model, criterion, optimizer, device,
                          dataloaders['train'], dataloaders['val'],
                          in_arg.epochs)

    # validation on test data
    print('\nGetting results on test data accuracy...')
    functions.test_model(model, criterion, device, dataloaders['test'])

    # saving checkpoint on trained model
    print('\nSaving checkpoint for current trained model...')
    functions.save_checkpoint(model, optimizer, dataset, in_arg.arch,
                              in_arg.epochs, in_arg.save_dir)
    print('Checkpoint saved!')
示例#2
0
def main():

    args = arg_parser_train()
    #Loading the data
    data_dir = args.data_dir
    train_dir = data_dir + '/train'
    valid_dir = data_dir + '/valid'
    test_dir = data_dir + '/test'
    image_datasets, dataloaders = train_transformer_load(
        train_dir, valid_dir, test_dir)

    # Label mapping
    import json
    with open('ImageClassifier/cat_to_name.json', 'r') as f:
        cat_to_name = json.load(f)

    # Load a pre_trained_network
    model, optimizer, criterion = pre_trained_network(args.arch, args.dropout,
                                                      args.hidden_units,
                                                      args.lr)

    learning(model, dataloaders['train'], dataloaders['valid'], args.epochs,
             40, criterion, optimizer, args.gpu)
    check_accuracy_on_test(dataloaders['test'], model)

    # TODO: Save the checkpoint
    save_checkpoint(model, args.save_dir, image_datasets['train'], args.arch,
                    args.hidden_units, args.lr, args.dropout)
示例#3
0
def main():
    parser = argparse.ArgumentParser(
        description='This program predicts a flower name from an image')
    parser.add_argument('data_dir', type=str, help='Dataset directory')
    parser.add_argument('--save_dir',
                        type=str,
                        default='./',
                        help='Saved checkpoint directory')
    parser.add_argument('--arch',
                        type=str,
                        default='vgg16',
                        help='Network architecture')
    parser.add_argument('--hidden_units',
                        type=int,
                        default='256',
                        help='Hidden units')
    parser.add_argument('--dropout',
                        type=float,
                        default='0.2',
                        help='Dropout for the hidden layers')
    parser.add_argument('--num_classes',
                        type=int,
                        default='256',
                        help='Number of classes for classification')
    parser.add_argument('--learning_rate',
                        type=float,
                        default='0.005',
                        help='Learning rate')
    parser.add_argument('--epochs',
                        type=int,
                        default='20',
                        help='Number of epochs')
    parser.add_argument('--gpu',
                        action='store_true',
                        help='Use GPU',
                        default=False)

    args = parser.parse_args()
    device = torch.device("cuda:0" if torch.cuda.is_available() else 'cpu'
                          ) if args.gpu else 'cpu'

    model, criterion, optimizer, scheduler = initialize_model(
        args.arch, args.hidden_units, args.dropout, args.num_classes, device,
        args.learning_rate)

    dataloaders, image_datasets = create_loaders(args.data_dir)

    train_model(model, dataloaders, criterion, optimizer, scheduler,
                image_datasets, args.epochs, device)
    test_model(model, dataloaders, image_datasets, device)
    save_checkpoint(model, optimizer, scheduler, args.epochs,
                    args.learning_rate, f'{args.arch}_checkpoint.pth')
示例#4
0
def main():

    # Get CLI arguments
    args = get_input_args()

    # Prep data
    train_transform = utilities.transform_data('train')
    test_transform = utilities.transform_data('test')

    # Dataloaders
    trainloader = utilities.load_data(args.data_directory + '/' + 'train',
                                      train_transform)
    validationloader = utilities.load_data(args.data_directory + '/' + 'valid',
                                           test_transform)

    # Setup and train model
    model, optimizer, criterion = functions.model_setup(
        args.arch, args.hidden_units, args.learning_rate)
    trained_model = functions.train_model(optimizer, criterion, model,
                                          trainloader, validationloader,
                                          args.gpu, args.epochs)

    # Save the model
    functions.save_checkpoint(trained_model, args.save_dir)
示例#5
0
                    dest="hidden_units",
                    default=256,
                    type=int)
parser.add_argument('--dropout', dest="dropout", default=0.2)
parser.add_argument('--learning_rate', dest="learning_rate", default=0.001)
parser.add_argument('--gpu', dest="gpu", default=True, type=bool)
parser.add_argument('--epochs', dest="epochs", default=2, type=int)
parser.add_argument('--print_every', dest="print_every", default=10, type=int)
parser.add_argument('--save_dir',
                    dest="save_dir",
                    default="./checkpoint.pth",
                    type=str)
parsed = parser.parse_args()
data_dir = parsed.data_dir
arch = parsed.arch
hidden_units = parsed.hidden_units
dropout = parsed.dropout
learning_rate = parsed.learning_rate
gpu = parsed.gpu
epochs = parsed.epochs
print_every = parsed.print_every
save_dir = parsed.save_dir

trainloader, validloader, class_to_idx = functions.load_data(data_dir)
model, criterion, optimizer = functions.setup(arch, hidden_units, dropout,
                                              learning_rate, gpu)
functions.train_network(model, criterion, optimizer, trainloader, validloader,
                        epochs, print_every, gpu)
functions.save_checkpoint(save_dir, arch, hidden_units, dropout, learning_rate,
                          gpu, model, optimizer, class_to_idx)
示例#6
0
                dest="arch",
                action="store",
                default="densenet121",
                type=str)
ap.add_argument('--hidden_units',
                type=int,
                dest="hidden_units",
                action="store",
                default=120)

pa = ap.parse_args()
path = pa.save_dir
lr = pa.learning_rate
structure = pa.arch
dropout = pa.dropout
hidden_layer1 = pa.hidden_units
power = pa.gpu
epochs = pa.epochs

trainloader, validationloader, testloader, train_data = functions.data()

model, optimizer, criterion = functions.setup(structure, dropout,
                                              hidden_layer1, lr, power)

functions.train_network(model, criterion, optimizer, trainloader,
                        validationloader, epochs, 20, power)

functions.save_checkpoint(train_data, epochs, structure, hidden_layer1,
                          dropout, lr, path)

print("The Model is trained")
    n_batches = len(loader_train)
    n_batches_test = len(loader_test)

    # Networks
    m = network.DQN(RESOLUTION, RESOLUTION, N_ACTIONS)
    m = m.to(DEVICE)
    o = optim.Adam(m.parameters(), lr=1e-5)

    start_epoch = 0
    run = f.Run(CHECKPOINT_DIR)
    start_epoch, m, o = f.load_checkpoint(run.get_checkpoint('32'), m, o)

    # validate_model(m, loader_test, idx_to_class_test)

    for epoch in range(N_EPOCHS):
        print('\n Epoch {}'.format(start_epoch + epoch))
        train_model(m, o, loader_train, idx_to_class_train)

        checkpoint = {
            'epoch': start_epoch + epoch + 1,
            'state_dict': m.state_dict(),
            'optimizer': o.state_dict()
        }
        f.save_checkpoint(checkpoint, CHECKPOINT_DIR, start_epoch + epoch)

        print('FIRST VALIDATION')
        validate_model(m, loader_test, idx_to_class_test, Q_TABLE_TEST,
                       "first")
        print('SECOND VALIDATION')
        validate_model(m, loader_test2, idx_to_class_test2, Q_TABLE_TEST2,
                       "second")
示例#8
0
else:
    logger = Logger(os.path.join(args.checkpoint, 'log.txt'), title=title)
    logger.set_names(['Learning Rate', 'Train Loss', 'Valid Loss','Train Acc.', 'Valid Acc.', 'Train Acc.5',
                      'Valid Acc.5'])

# Train and validate
for epoch in range(start_epoch, args.epochs):
    state['lr'] = adjust_learning_rate(state['lr'], optimizer, epoch, args.gamma, args.schedule)
    print('\nEpoch: [%d | %d] LR: %f' % (epoch + 1, args.epochs, state['lr']))

    train_loss, train_acc, train_acc5 = train_one_epoch(trainloader, model, criterion, optimizer, use_cuda=use_cuda)
    test_loss, test_acc, test_acc5 = test(testloader, model, criterion, use_cuda=use_cuda)
    # append logger file
    logger.append([state['lr'], train_loss, test_loss, train_acc, test_acc, train_acc5, test_acc5])

    # save model ap
    is_best = test_acc > best_acc
    best_acc = max(test_acc, best_acc)
    if do_save_checkpoint:
        save_checkpoint({
                'epoch': epoch + 1,
                'state_dict': model.state_dict(),
                'acc': test_acc,
                'best_acc': best_acc,
                'optimizer' : optimizer.state_dict(),
            }, is_best, checkpoint=args.checkpoint)

logger.close()
print('Best acc:', best_acc)

示例#9
0
def main():
    #define supported models
    allowed_models = [
        'ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101', 'resnet152'
    ]

    #Set up argument parser for console input
    parser = argparse.ArgumentParser(description='Train NN')
    parser.add_argument('data_dir',
                        help='directory containing sub-folders with data')
    parser.add_argument('--save_dir',
                        help='directory for saving checkpoint',
                        default='checkpoints')
    parser.add_argument('--arch',
                        help='pre-trained model architecture',
                        default='resnet18',
                        choices=allowed_models)
    parser.add_argument('--learning_rate',
                        help='learning rate during learning',
                        type=float,
                        default=0.01)
    parser.add_argument('--dropout',
                        help='dropout during learning',
                        type=float,
                        default=0.05)
    parser.add_argument('--hidden_units',
                        help='List of number of nodes in hidden layers',
                        nargs='+',
                        type=int,
                        default=[256, 128])
    parser.add_argument('--epochs',
                        help='Number of epochs for training',
                        default=3,
                        type=int)
    parser.add_argument('--gpu', help='Enable GPU', action='store_true')

    args = parser.parse_args()

    # Describe directories relative to working directory
    data_dir = args.data_dir
    train_dir = data_dir + '/train'
    valid_dir = data_dir + '/valid'
    test_dir = data_dir + '/test'
    save_dir = args.save_dir

    # Set variables for console input arguments
    model_arch = args.arch
    model_hidden_units = args.hidden_units
    learning_rate = args.learning_rate
    drop = args.dropout

    #Testing area
    print('Data directory: ' + data_dir)
    print('hidden units: ' + str(args.hidden_units))
    print('Save directory: ' + save_dir)
    print('Architecture: ' + args.arch)

    #create save directory if not existing
    fu.create_directory(save_dir)

    # Loading Pre-Trained model dependent on console input arch
    model = models.__getattribute__(model_arch)(pretrained=True)

    # Freeze parameters so we don't backprop through them
    for param in model.parameters():
        param.requires_grad = False

    # Create the network, define the criterion and optimizer
    model.fc = fu.Network(model.fc.in_features, 102, model_hidden_units, drop)
    criterion = nn.NLLLoss()
    optimizer = optim.Adam(model.fc.parameters(), lr=learning_rate)

    device = torch.device(
        'cuda' if torch.cuda.is_available() and args.gpu == True else 'cpu')
    print('Device is: ', device)

    epochs = args.epochs
    print_every = 50
    running_loss = 0
    steps = 0

    train_loader, test_loader, valid_loader, train_data, test_data, valid_data = load_transform.load_transform(
        data_dir, train_dir, valid_dir, test_dir)

    fu.train(device, model, epochs, criterion, optimizer, print_every,
             train_loader, test_loader, valid_loader)
    fu.save_checkpoint(model, model_arch, epochs, criterion, optimizer,
                       train_data, save_dir)

    return model, test_loader, criterion
def main(args):
    logging.info(args)
    writer = SummaryWriter(args.save_dir)

    best_prec1, best_ata, cln_best_prec1 = 0, 0, 0
    if not torch.cuda.is_available():
        logging.info('no gpu device available')
        sys.exit(1)
    cudnn.benchmark = True

    setup_seed(args.seed)

    model = get_network_func(args.norm_module)
    model.cuda()

    start_epoch = 0
    if args.resume:
        print('resume from checkpoint')
        checkpoint = torch.load(os.path.join(args.save_dir, 'model.pt'))
        start_epoch = checkpoint['epoch']
        model.load_state_dict(checkpoint['state_dict'])

    # dataset process
    train_datasets = datasets.CIFAR10(root=args.data,
                                      train=True,
                                      transform=transforms.Compose([
                                          transforms.RandomHorizontalFlip(),
                                          transforms.RandomCrop(32, 4),
                                          transforms.ToTensor()
                                      ]),
                                      download=True)

    test_dataset = datasets.CIFAR10(root=args.data,
                                    train=False,
                                    transform=transforms.Compose(
                                        [transforms.ToTensor()]),
                                    download=True)

    test_loader = torch.utils.data.DataLoader(test_dataset,
                                              batch_size=args.batch_size,
                                              shuffle=False,
                                              num_workers=2,
                                              pin_memory=True)

    valid_size = 0.1
    indices = list(range(len(train_datasets)))
    split = int(np.floor(valid_size * len(train_datasets)))

    np.random.seed(args.seed)
    np.random.shuffle(indices)

    train_idx, valid_idx = indices[split:], indices[:split]
    train_sampler = SubsetRandomSampler(train_idx)
    valid_sampler = SubsetRandomSampler(valid_idx)

    train_loader = torch.utils.data.DataLoader(train_datasets,
                                               batch_size=args.batch_size,
                                               sampler=train_sampler)

    val_loader = torch.utils.data.DataLoader(train_datasets,
                                             batch_size=args.batch_size,
                                             sampler=valid_sampler)

    decreasing_lr = list(map(int, args.decreasing_lr.split(',')))

    criterion = nn.CrossEntropyLoss()
    criterion = criterion.cuda()

    optimizer = torch.optim.SGD(model.parameters(),
                                args.lr,
                                momentum=args.momentum,
                                weight_decay=args.weight_decay)

    scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer,
                                                     milestones=decreasing_lr,
                                                     gamma=0.1)

    print('starting adv training')

    all_result = {}
    train_acc = []
    ta0 = []
    ata0 = []
    ta1 = []
    ata1 = []

    test_ta0 = []
    test_ata0 = []
    test_ta1 = []
    test_ata1 = []

    if os.path.exists(args.save_dir) is not True:
        os.mkdir(args.save_dir)

    for epoch in tqdm(range(args.epochs)):

        if epoch < start_epoch:
            scheduler.step()
            continue

        print(optimizer.state_dict()['param_groups'][0]['lr'])
        acc, train_loss = train(args, train_loader, model, criterion,
                                optimizer, epoch)
        writer.add_scalar('acc/train_acc', acc, epoch)
        writer.add_scalar('loss/train_loss', train_loss, epoch)

        # evaluate on validation set
        tacc_clean, tloss_clean = validate(args, val_loader, model, criterion,
                                           0)
        writer.add_scalar('acc/tacc_clean', tacc_clean, epoch)
        writer.add_scalar('loss/tloss_clean', tloss_clean, epoch)

        atacc_clean, atloss_clean = validate_adv(args, val_loader, model,
                                                 criterion, 0)
        writer.add_scalar('acc/atacc_clean', atacc_clean, epoch)
        writer.add_scalar('loss/atloss_clean', atloss_clean, epoch)

        tacc_adv, tloss_adv = validate(args, val_loader, model, criterion, 1)
        writer.add_scalar('acc/tacc_adv', tacc_adv, epoch)
        writer.add_scalar('loss/tloss_adv', tloss_adv, epoch)

        atacc_adv, atloss_adv = validate_adv(args, val_loader, model,
                                             criterion, 1)
        writer.add_scalar('acc/atacc_adv', atacc_adv, epoch)
        writer.add_scalar('loss/atloss_adv', atloss_adv, epoch)

        # evaluate on test set
        # clean branch
        test_tacc_clean, test_tloss_clean = validate(args, test_loader, model,
                                                     criterion, 0)
        writer.add_scalar('acc/test_tacc_clean', test_tacc_clean, epoch)
        writer.add_scalar('loss/test_tloss_clean', test_tloss_clean, epoch)

        test_atacc_clean, test_atloss_clean = validate_adv(
            args, test_loader, model, criterion, 0)
        writer.add_scalar('acc/test_atacc_clean', test_atacc_clean, epoch)
        writer.add_scalar('loss/test_atloss_clean', test_atloss_clean, epoch)

        # adv branch
        test_tacc_adv, test_tloss_adv = validate(args, test_loader, model,
                                                 criterion, 1)
        writer.add_scalar('acc/test_tacc_adv', test_tacc_adv, epoch)
        writer.add_scalar('loss/test_tloss_adv', test_tloss_adv, epoch)

        test_atacc_adv, test_atloss_adv = validate_adv(args, test_loader,
                                                       model, criterion, 1)
        writer.add_scalar('acc/test_atacc_adv', test_atacc_adv, epoch)
        writer.add_scalar('loss/test_atloss_adv', test_atloss_adv, epoch)

        scheduler.step()

        train_acc.append(acc)
        ta0.append(tacc_clean)
        ata0.append(atacc_clean)
        ta1.append(tacc_adv)
        ata1.append(atacc_adv)

        test_ta0.append(test_tacc_clean)
        test_ata0.append(test_atacc_clean)
        test_ta1.append(test_tacc_adv)
        test_ata1.append(test_atacc_adv)

        # remember best prec@1 and save checkpoint
        is_best = tacc_adv > best_prec1
        best_prec1 = max(tacc_adv, best_prec1)

        ata_is_best = atacc_adv > best_ata
        best_ata = max(atacc_adv, best_ata)

        cln_is_best = tacc_clean > cln_best_prec1
        cln_best_prec1 = max(tacc_clean, cln_best_prec1)

        if is_best:
            save_checkpoint(
                {
                    'epoch': epoch + 1,
                    'state_dict': model.state_dict(),
                    'best_prec1': best_prec1,
                    'ata_best_prec1': best_ata,
                },
                is_best,
                filename=os.path.join(args.save_dir, 'best_model.pt'))

        if cln_is_best:
            save_checkpoint(
                {
                    'epoch': epoch + 1,
                    'state_dict': model.state_dict(),
                    'best_prec1': best_prec1,
                    'ata_best_prec1': best_ata,
                },
                is_best,
                filename=os.path.join(args.save_dir, 'clean_best_model.pt'))

        if ata_is_best:
            save_checkpoint(
                {
                    'epoch': epoch + 1,
                    'state_dict': model.state_dict(),
                    'best_prec1': best_prec1,
                    'ata_best_prec1': best_ata,
                },
                is_best,
                filename=os.path.join(args.save_dir, 'ata_best_model.pt'))

        save_checkpoint(
            {
                'epoch': epoch + 1,
                'state_dict': model.state_dict(),
                'best_prec1': best_prec1,
                'ata_best_prec1': best_ata,
            },
            is_best,
            filename=os.path.join(args.save_dir, 'model.pt'))

        if epoch and epoch % 10 == 0:
            save_checkpoint(
                {
                    'epoch': epoch + 1,
                    'state_dict': model.state_dict(),
                    'best_prec1': best_prec1,
                    'ata_best_prec1': best_ata,
                },
                is_best,
                filename=os.path.join(args.save_dir, f'model_epoch{epoch}.pt'))

        all_result['train'] = train_acc
        all_result['test_ta0'] = test_ta0
        all_result['test_ata0'] = test_ata0
        all_result['test_ta1'] = test_ta1
        all_result['test_ata1'] = test_ata1
        all_result['ta0'] = ta0
        all_result['ata0'] = ata0
        all_result['ta1'] = ta1
        all_result['ata1'] = ata1

        pickle.dump(all_result,
                    open(os.path.join(args.save_dir, 'result_c.pkl'), 'wb'))
示例#11
0
                    dest='dropout',
                    action='store',
                    type=int,
                    default=0.05)
# parser.add_argument('--lr' , dest='lr' , action ='store' , type = int, default= 0.001)
parser.add_argument('--save_directory',
                    dest='save_directory',
                    action='store',
                    default='./checkpoint.pth')

parser = parser.parse_args()
epochs = parser.epochs
lr = parser.lr
structure = parser.structure
dropout = parser.dropout
hidden_layer1 = parser.hidden_layer1
hidden_layer2 = parser.hidden_layer2
power = parser.gpu

train_loaders, valid_loaders, test_loaders = functions.read_data(data_dir)
## load the model
model, optimizer, criterion = functions.model_setup(structure, epochs, dropout,
                                                    hidden_layer1,
                                                    hidden_layer2, lr, power)
## train the model
functions.train(model, epochs, criterion, optimizer, train_loaders,
                valid_loaders, power)
## save the model
functions.save_checkpoint(path, hidden_layer1, hidden_layer2, dropout, lr,
                          epochs)
示例#12
0
# 5. Build Custom Classifier

from collections import OrderedDict

classifier = nn.Sequential(
    OrderedDict([('fc1', nn.Linear(25088, 5000)), ('relu', nn.ReLU()),
                 ('drop', nn.Dropout(p=0.5)), ('fc2', nn.Linear(5000, 3)),
                 ('output', nn.LogSoftmax(dim=1))]))
model.classifier = classifier

# 6. Loss function and gradient descent

criterion = nn.NLLLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 7. model training

model = functions.train_classifier(model, train_loader, validate_loader,
                                   optimizer, criterion)

# 8. test accuracy

functions.test_accuracy(model, test_loader)

# 9. model save

model_save_path = r"C:/Users/LG/Desktop/ksb/3. CODE/model/"
filename = 'sepa_image_classifier2.pth'
functions.save_checkpoint(model, training_dataset, model_save_path, filename)
示例#13
0
parser.add_argument('data_dir', action="store")

parser.add_argument('--save_dir', action="store", dest="save_dir", default = "checkpoint.pth")
parser.add_argument('--arch', action="store", dest="arch", default = "densenet121")
parser.add_argument('--learning_rate', action="store", dest="learning_rate", type=int, default = 0.003)
parser.add_argument('--hidden_units', action="store", dest="hidden_units", type=int, default = 1024)
parser.add_argument('--epochs', action="store", dest="epochs", type=int, default = 2)
parser.add_argument('--gpu', action="store", dest="gpu", default = "gpu")

parse_args = parser.parse_args()

data_dir = parse_args.data_dir

save_path = parse_args.save_dir
arch = parse_args.arch
learning_rate = parse_args.learning_rate
hidden_units = parse_args.hidden_units
epochs = parse_args.epochs
gpu = parse_args.gpu


train_datasets, valid_datasets, test_datasets, trainloader, validloader, testloader =  data_load(data_dir)

model, device, output_units, criterion, optimizer = define_model(arch, learning_rate, hidden_units, gpu)

model_train(trainloader, validloader, testloader, epochs, model, device, output_units, criterion, optimizer)

save_checkpoint(save_path, arch, hidden_units, learning_rate, epochs)


hidden_unit = results.hidden_unit
class_number = results.class_number
structure = results.structure
gpu_mode = results.gpu

# load and preprocess data
image_datasets_train, image_datasets_valid, image_datasets_test, dataloaders_train, dataloaders_valid, dataloaders_test = load_data(
    data_dir)

# load pretrained model and add new classifiers
model = getattr(models, arch)(pretrained=True)
build_classifier(model, input_unit, hidden_unit, class_number, dropout)

# set criterion and optimizer
criterion = nn.NLLLoss()
optimizer = optim.Adam(model.classifier.parameters(), lr=learning_rate)

# model training and validation
print('prerequisite has been fulfiled, now start training & validation...')
model = train_and_validation_the_model(gpu_mode, model, learning_rate, epochs,
                                       dataloaders_train, dataloaders_valid,
                                       criterion, optimizer)

# model testing
print('\n now start testing...')
test_model(model, dataloaders_test, gpu_mode)

# model saving
print('\n now start saving...')
checkpoint = save_checkpoint(model, image_datasets_train, chekpoint_name, arch)
arguments = parser.parse_args()

data_dir = arguments.data_dir
save_dir = arguments.save_dir
arch = arguments.arch
learning_rate = arguments.learning_rate
hidden_units = arguments.hidden_units
epochs = arguments.epochs
gpu = arguments.gpu

print("Data Info :\n") 
train_dataloader, test_dataloader, validation_dataloader = load_data(data_dir)
print("\nModel Info :\n") 
model = create_model(arch, hidden_units)
print("\nTrain starts:\n")

if gpu:
    print("GPU activated")
else: 
    print("GPU is not activated")

print("Training epochs : {}".format(str(epochs)))
print("Learning rate : {:.4f}".format(learning_rate))
      
model, optimizer, criterion = train(model, train_dataloader, validation_dataloader, learning_rate, gpu, epochs, 40)
print("\nSave checkpoint:\n")
save_checkpoint(model, save_dir, optimizer, criterion, epochs)