Function main

Language_Model/run_exp.py:95–214

(args)

def main(args):
    # Seed the experiment, for repeatability
    seed_experiment(args.seed)
    # Dataloaders
    train_dataset = Wikitext2(args.data_folder, split="train")
    train_dataloader = DataLoader(
        train_dataset,
        batch_size=args.batch_size,
        shuffle=True,
        num_workers=args.num_workers,
    )
    valid_dataset = Wikitext2(args.data_folder, split="validation")
    valid_dataloader = DataLoader(
        valid_dataset,
        batch_size=args.batch_size,
        shuffle=False,
        num_workers=args.num_workers,
    )
    test_dataset = Wikitext2(args.data_folder, split="test")
    test_dataloader = DataLoader(
        test_dataset,
        batch_size=args.batch_size,
        shuffle=False,
        num_workers=args.num_workers,
    )

    # Download the embeddings
    if not os.path.isfile(args.embeddings):
        print("Downloading embeddings...")
        urllib.request.urlretrieve(EMBEDDINGS_URL, args.embeddings)

    # Model

    model = MiniGPT1.load_embeddings_from(
        args.embeddings, num_heads=12, num_layers=args.layers
    )
    model.to(args.device)

    # Optimizer
    if args.optimizer == "adamw":
        optimizer = AdamW(
            model.parameters(), lr=args.lr, weight_decay=args.weight_decay
        )
    elif args.optimizer == "AdaFisherW":
        optimizer = AdaFisherW(model, lr=args.lr, gammas=[args.gamma1, args.gamma2], TCov=args.curvature_update_interval,
                               weight_decay=args.weight_decay, Lambda=args.damping
                               )
    elif args.optimizer == "AdaHessian":
        optimizer = Adahessian(
            model.parameters(),
            lr=args.lr,
            weight_decay=args.weight_decay,
        )
    elif args.optimizer in ['Shampoo', 'kfac']:
        optimizer = SGD(
            model.parameters(),
            lr=args.lr,
            weight_decay=args.weight_decay,