!pip install accelerate

"""
Explanation: Hugging Face Accelerate Demo
Note: Before running this demo, please make sure that you have wandb.ai free account. 
Let us install accelerate.
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

# This is a demo of the PyTorch Accelerate API.
from accelerate import Accelerator

"""
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="accelerate-project", entity="upeee", config=config)

"""
Explanation: wandb initialization. See wandb_demo notebook for more details.
End of explanation
"""


# Shows the code to be replaced with the Accelerate API.
#device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
accelerator = Accelerator()

model = torchvision.models.resnet18(pretrained=False, progress=True)

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

# Replace the model with the Accelerate API.
#model.to(device)

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

"""
Explanation: Build the model
Use a ResNet18 from torchvision. See wandb_demo notebook for more details.
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)

# Accelerate API
model = accelerator.prepare(model)
optimizer = accelerator.prepare(optimizer)
scheduler = accelerator.prepare(scheduler)
train_loader = accelerator.prepare(train_loader)


"""
Explanation: Loss function, Optimizer, Scheduler and DataLoader
See wandb_demo notebook for more details.
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()
test_loader = accelerator.prepare(test_loader)
image = image.to(accelerator.device)
model.eval()
with torch.no_grad():
  pred = torch.argmax(model(image), 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: Visulaizing sample data from test split
See wandb_demo notebook for more details.
Note the last line that uses Accelerate API to wrap the model, optimizer, data loaders and scheduler.
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()
    # Replaced by the Accelerate API.
    #target = target.to(device)
    #output = model(data.to(device))
    
    output = model(data)
    loss_value = loss(output, target)

    # Replaced by the Accelerate API.
    #loss_value.backward()
    accelerator.backward(loss_value)

    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
Using Accelerate, we do not need to transfer the model to the device.
See wandb_demo notebook for more details.
End of explanation
"""


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

      # Replaced by the Accelerate API.
      #output = model(data.to(device))   
      #target = target.to(device)

      output = model(data)
      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. Again, no need to transfer the data to the device.
See wandb_demo notebook for more details.
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
        accelerator.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), 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. 
See wandb_demo notebook for more details.
End of explanation
"""


model = torch.load("resnet18_best_acc.pth")
# Using Accelerator API
model = accelerator.prepare(model)
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
"""


