text: import torch
  from tqdm import tqdm
  import torch.nn as nn
  import torch.optim as optim
  import torchvision.transforms as transforms
  from torch.utils.tensorboard import SummaryWriter # For TensorBoard
  from utils import save_checkpoint, load_checkpoint, print_examples
  from dataset import get_loader
  from model import SeqToSeq
  from tabulate import tabulate # To tabulate loss and epoch
  import argparse
  import json
  
  def main(args):
      transform = transforms.Compose(
          [
              transforms.Resize((356, 356)),
              transforms.RandomCrop((299, 299)),
              transforms.ToTensor(),
              transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
          ]
      )
  
      train_loader, _ = get_loader(
          root_folder = args.root_dir,
          annotation_file = args.csv_file,
          transform=transform,
          batch_size = 64,
          num_workers=2,
      )
      vocab = json.load(open('vocab.json'))
  
      torch.backends.cudnn.benchmark = True
      device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
      load_model = False
      save_model = True
      train_CNN = False
  
      # Hyperparameters
      embed_size = args.embed_size
      hidden_size = args.hidden_size
      vocab_size = len(vocab['stoi'])
      num_layers = args.num_layers
      learning_rate = args.lr
      num_epochs = args.num_epochs
      # for tensorboard
  
      
      writer = SummaryWriter(args.log_dir)
      step = 0
      model_params = {'embed_size': embed_size, 'hidden_size': hidden_size, 'vocab_size':vocab_size, 'num_layers':num_layers}
      # initialize model, loss etc
      model = SeqToSeq(**model_params, device = device).to(device)
      criterion = nn.CrossEntropyLoss(ignore_index = vocab['stoi']["<PAD>"])
      optimizer = optim.Adam(model.parameters(), lr=learning_rate)
  
      # Only finetune the CNN
      for name, param in model.encoder.inception.named_parameters():
          if "fc.weight" in name or "fc.bias" in name:
              param.requires_grad = True
          else:
              param.requires_grad = train_CNN
  
      #load from a save checkpoint
      if load_model:
          step = load_checkpoint(torch.load(args.save_path), model, optimizer)
  
      model.train()
      best_loss, best_epoch = 10, 0
      for epoch in range(num_epochs):
          print_examples(model, device, vocab['itos'])
  
          for idx, (imgs, captions) in tqdm(
              enumerate(train_loader), total=len(train_loader), leave=False):
              imgs = imgs.to(device)
              captions = captions.to(device)
  
              outputs = model(imgs, captions[:-1])
              loss = criterion(
                  outputs.reshape(-1, outputs.shape[2]), captions.reshape(-1)
              )
  
              writer.add_scalar("Training loss", loss.item(), global_step=step)
              step += 1
  
              optimizer.zero_grad()
              loss.backward(loss)
              optimizer.step()
  		
          train_loss = loss.item()
          if train_loss < best_loss:
              best_loss = train_loss
              best_epoch = epoch + 1
              if save_model:
                  checkpoint = {
                      "model_params": model_params,
  		            "state_dict": model.state_dict(),
  		            "optimizer": optimizer.state_dict(),
  		            "step": step
  		        }
                  save_checkpoint(checkpoint, args.save_path)
  
  
          table = [["Loss:", train_loss],
  				["Step:", step],
                  ["Epoch:", epoch + 1],
  		 		["Best Loss:", best_loss],
  		  		["Best Epoch:", best_epoch]]
          print(tabulate(table))
  	
  	
  if __name__ == "__main__":
  
      parser = argparse.ArgumentParser()
  
      parser.add_argument('--root_dir', type = str, default = './flickr30k/flickr30k_images', help = 'path to images folder')
      parser.add_argument('--csv_file', type = str, default = './flickr30k/results.csv', help = 'path to captions csv file')
      parser.add_argument('--log_dir', type = str, default = './drive/MyDrive/TensorBoard/', help = 'path to save tensorboard logs')
      parser.add_argument('--save_path', type = str, default = './drive/MyDrive/checkpoints/Seq2Seq.pt', help = 'path to save checkpoint')
      # Model Params
      parser.add_argument('--batch_size', type = int, default = 64)
      parser.add_argument('--num_epochs', type = int, default = 100)
      parser.add_argument('--embed_size', type = int, default=256)
      parser.add_argument('--hidden_size', type = int, default=512)
      parser.add_argument('--lr', type = float, default= 0.001)
      parser.add_argument('--num_layers', type = int, default = 3, help = 'number of lstm layers')
  
      args = parser.parse_args()
      
      main(args)