!pip install wandb

"""
Explanation: Weights and Biases (wandb) Demo
In deep learning, we perform a lot of model training especially for novel neural architectures. The problem is deep learning frameworks like PyTorch do not provide sufficient tools to visualize input data, track the progress of our experiments, log data, and visualize the outputs. 
wandb addresses this problem. In this demo, we will train a ResNet18 model from scratch. We show how to use wandb to visualize input data, prediction, and training progress using loss function value and accuracy. 
Note: Before running this demo, please make sure that you have wandb.ai free account. 
Let us install wandb.
End of explanation
"""


import torch
import torchvision
import wandb
import datetime
from torch.optim import SGD
from torch.optim.lr_scheduler import CosineAnnealingLR
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from ui import progress_bar

"""
Explanation: Import the required modules.
End of explanation
"""


wandb.login()
config = {
  "learning_rate": 0.1,
  "epochs": 100,
  "batch_size": 128,
  "dataset": "cifar10"
}
run = wandb.init(project="wandb-project", entity="upeee", config=config)

"""
Explanation: Login to and initialize wandb. You will need to use your wandb API key to run this demo.
As the config indicates, we will train our model using cifar10 dataset, learning rate of 0.1, and batch size of 128 for 100 epochs. 
epochs means a complete sampling of the dataset (train). In the wandb plots, step is the term used instead of epoch.
batch size is the number of samples per training step.
End of explanation
"""


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = torchvision.models.resnet18(pretrained=False, progress=True)

model.fc = torch.nn.Linear(model.fc.in_features, 10)  
model.to(device)

# watch model gradients during training
wandb.watch(model)

"""
Explanation: Build the model
Use a ResNet18 from torchvision. Remove the last layer that was used for 1k-class ImageNet classification. Since we will use CIFAR10, the last layer is replaced by a linear layer with 10 outputs. We will train the model from scratch, so we set pretrained=False.
End of explanation
"""


loss = torch.nn.CrossEntropyLoss()
optimizer = SGD(model.parameters(), lr=wandb.config.learning_rate)
scheduler = CosineAnnealingLR(optimizer, T_max=wandb.config.epochs)

x_train = datasets.CIFAR10(root='./data', train=True, 
                           download=True, 
                           transform=transforms.ToTensor())
x_test = datasets.CIFAR10(root='./data',
                          train=False, 
                          download=True, 
                          transform=transforms.ToTensor())
train_loader = DataLoader(x_train, 
                          batch_size=wandb.config.batch_size, 
                          shuffle=True, 
                          num_workers=2)
test_loader = DataLoader(x_test, 
                         batch_size=wandb.config.batch_size, 
                         shuffle=False, 
                         num_workers=2)

"""
Explanation: Loss function, Optimizer, Scheduler and DataLoader
The appropriate loss function is cross entropy for multi-category classfications. We use SGD or stochastic gradient descent for optimization. Our learning rate that starts at 0.1 decays to zero at the end of total number of epochs. The decay is controlled by a cosine learning rate decay scheduler. 
Finally, we use cifar10 dataset that is available in torchvision. We will discuss datasets and dataloaders in our future demo. For the meantime, we can treat dataloader as a data strcuture that dispenses batch size data from either the train or test split of the dataset.
End of explanation
"""



label_human = ["airplane", "automobile", "bird", "cat", "deer", "dog", "frog", "horse", "ship", "truck"]

table_test = wandb.Table(columns=['Image', "Ground Truth", "Initial Pred Label",])

image, label = iter(test_loader).next()
model.eval()
with torch.no_grad():
  pred = torch.argmax(model(image.to(device)), dim=1).cpu().numpy()

for i in range(8):
  table_test.add_data(wandb.Image(image[i]),
                      label_human[label[i]], 
                      label_human[pred[i]])
  print(label_human[label[i]], "vs. ",  label_human[pred[i]])

"""
Explanation: Visualizing sample data from test split
We can visualize data from the test split by getting a batch sample: image, label = iter(test_loader).next(). We use wandb table to create a column for image, grount truth label and initial model predicted label. The wandb table will show up when we run wandb.log() during training. 
CIFAR10 dataset is made of small 32x32 RGB images. Each image belongs to one of the 10 categories or classes. Below are sample images from CIFAR10 and their corresponding human labels.
<img src="cifar10-samples.png" width="600" height="600">
End of explanation
"""


def train(epoch):
  model.train()
  train_loss = 0
  correct = 0
  train_samples = 0

  # sample a batch. compute loss and backpropagate
  for batch_idx, (data, target) in enumerate(train_loader):
    optimizer.zero_grad()
    target = target.to(device)
    output = model(data.to(device))
    loss_value = loss(output, target)
    loss_value.backward()
    optimizer.step()
    scheduler.step(epoch)
    train_loss += loss_value.item()
    train_samples += len(data)
    pred = output.argmax(dim=1, keepdim=True)
    correct += pred.eq(target.view_as(pred)).sum().item()
    if batch_idx % 10 == 0:
      accuracy = 100. * correct / len(train_loader.dataset)
      progress_bar(batch_idx,
                   len(train_loader),
                  'Train Epoch: {}, Loss: {:.6f}, Acc: {:.2f}%'.format(epoch+1, 
                  train_loss/train_samples, accuracy))
  
  train_loss /= len(train_loader.dataset)
  accuracy = 100. * correct / len(train_loader.dataset)

  return accuracy, train_loss

"""
Explanation: The train loop
At every epoch, we will run the train loop for the model. At every iteration, we will get a batch of data from the train split. We will use the data to update the model parameters. We will use the loss function to calculate the loss value. We will use the optimizer to update the model parameters. We will use the scheduler to update the learning rate. Later, we will use the wandb table to visualize the loss and accuracy.
We use progress_bar to show the progress of the training.
End of explanation
"""


def test():
  model.eval()
  test_loss = 0
  correct = 0
  with torch.no_grad():
    for data, target in test_loader:
      output = model(data.to(device))   
      target = target.to(device)

      test_loss += loss(output, target).item()
      pred = output.argmax(dim=1, keepdim=True)
      correct += pred.eq(target.view_as(pred)).sum().item()

  test_loss /= len(test_loader.dataset)
  accuracy = 100. * correct / len(test_loader.dataset)

  print('\nTest Loss: {:.4f}, Acc: {:.2f}%\n'.format(test_loss, accuracy))

  return accuracy, test_loss

"""
Explanation: The validation loop
After every epoch, we will run the validation loop for the model. In this way, we can track the progress of our model training. Both the average loss and accuracy are calculated. During training, we will use the wandb table to visualize the loss and accuracy.
End of explanation
"""


run.display(height=1000)

start_time = datetime.datetime.now()
best_acc = 0
for epoch in range(wandb.config["epochs"]):
    train_acc, train_loss = train(epoch)
    test_acc, test_loss = test()
    if test_acc > best_acc:
        wandb.run.summary["Best accuracy"] = test_acc
        best_acc = test_acc
        torch.save(model, "resnet18_best_acc.pth")
    wandb.log({
        "Train accuracy": train_acc,
        "Test accuracy": test_acc,
        "Train loss": train_loss,
        "Test loss": test_loss,
        "Learning rate": optimizer.param_groups[0]['lr']
    })

elapsed_time = datetime.datetime.now() - start_time
print("Elapsed time: %s" % elapsed_time)
wandb.run.summary["Elapsed train time"] = str(elapsed_time)

model.eval()
with torch.no_grad():
  pred = torch.argmax(model(image.to(device)), dim=1).cpu().numpy()

final_pred = []
for i in range(8):
    final_pred.append(label_human[pred[i]])
    print(label_human[label[i]], "vs. ",  final_pred[i])

table_test.add_column(name="Final Pred Label", data=final_pred)

wandb.log({"Test data": table_test})

wandb.finish()

"""
Explanation: wandb plots
Finally, we will use wandb to visualize the training progress. We will use the following plots:
- Model gradients (wandb.watch(model))
- Train and test losses ("train loss": train_loss, "test loss": test_loss,)
- Train and validation accuracies ("Train accuracy": train_acc, "Test accuracy": test_acc,)
- Learning rate which decreases over epochs ("Learning rate": optimizer.param_groups[0]['lr'])
We re-use the earlier table_test to see the final prediction.
We save the best peforming model to ./resnet18_best_acc.pth. This can be used as a pretrained model like the pre-trained model in torchvision.
End of explanation
"""


model = torch.load("resnet18_best_acc.pth")
accuracy, _ = test()
print("Best accuracy: %.2f" % accuracy)

"""
Explanation: Load the best performing model
In the following code, we load the best performing model. The model is saved in ./resnet18_best_acc.pth. The average accuracy of the model is the same as the one in the previous section.
End of explanation
"""


