Function train

sweep/cnn/train.py:15–119

()

from config import hyperparameter_defaults

def train():
    wandb.init(config=hyperparameter_defaults)
    config = wandb.config
    transform = transforms.Compose([transforms.ToTensor(),
                                    transforms.Normalize((0.1307,), (0.3081,))])

    train_dataset = fashion(root='./data',
                                train=True,
                                transform=transform,
                                download=True
                               )

    test_dataset = fashion(root='./data',
                                train=False,
                                transform=transform,
                               )

    label_names = [
        "T-shirt or top",
        "Trouser",
        "Pullover",
        "Dress",
        "Coat",
        "Sandal",
        "Shirt",
        "Sneaker",
        "Bag",
        "Boot"]

    train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
                                               batch_size=config.batch_size,
                                               shuffle=True)

    test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
                                              batch_size=config.batch_size,
                                              shuffle=False)


    model = CNNModel(config)
    wandb.watch(model)

    criterion = nn.CrossEntropyLoss()

    optimizer = torch.optim.Adam(model.parameters(), lr=config.learning_rate)

    iter = 0
    for epoch in range(config.epochs):
        for i, (images, labels) in enumerate(train_loader):

            images = Variable(images)
            labels = Variable(labels)

            # Clear gradients w.r.t. parameters
            optimizer.zero_grad()

            # Forward pass to get output/logits
            outputs = model(images)