text: #!/usr/bin/env python3
  import os
  import pickle
  import argparse
  
  import numpy as np
  import torch
  import torch.nn as nn
  
  from PIL import Image
  from tqdm import tqdm
  
  from models.OracleResnetModel import OracleResnet
  from data.faceattribute_dataset import FaceAttributesDataset
  
  device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
  
  parser = argparse.ArgumentParser()
  parser.add_argument('--checkpoint_dir', type=str, default='/path/to/checkpoints', help='the path of the checkpoints')
  parser.add_argument('--dataroot', type=str, default='/path/to/dataroot', help='the path of the dataroot')
  parser.add_argument('--celeba_or_celebamhq', type=str, default='celebamhq')
  config=parser.parse_args()
  
  
  class Args:
  
      if config.celeba_or_celebamhq == "celebamhq":
          oracle_name = 'celebamaskhq'
          image_path_train = os.path.join(config.dataroot, "CelebAMask-HQ", "CelebAMask-HQ", "train", "images")
          image_path_val = os.path.join(config.dataroot, "CelebAMask-HQ", "CelebAMask-HQ", "test", "images")
          attributes_path = os.path.join(config.dataroot, "CelebAMask-HQ", "CelebAMask-HQ", "CelebAMask-HQ-attribute-anno.txt")
  
      elif config.celeba_or_celebamhq == "celeba":
          oracle_name = 'celeba'
          image_path_train = os.path.join(config.dataroot, "celeba_squared_128", "img_squared128_celeba_train")
          image_path_val = os.path.join(config.dataroot, "celeba_squared_128", "img_squared128_celeba_test")
          attributes_path = os.path.join(config.dataroot, "celeba_squared_128", "list_attr_celeba.txt")
  
      optimizer = 'adam'
      lr = 0.0001
      step_size = 10
      gamma_scheduler = 0.5
      num_epochs = 5
  
      oracle_pretraining_path = os.path.join(config.checkpoint_dir, "vggface2_pretrainings_for_oracle/resnet50_ft_dag.pth")
  
  opt=Args()
  
  # Prepare data
  data_train = FaceAttributesDataset(image_path=opt.image_path_train,attributes_path=opt.attributes_path,load_size=(224,224))
  data_val = FaceAttributesDataset(image_path=opt.image_path_val,attributes_path=opt.attributes_path,load_size=(224,224))
  
  dataloader_train = torch.utils.data.DataLoader(data_train,batch_size=32, shuffle=True,num_workers=4)
  dataloader_val = torch.utils.data.DataLoader(data_val,batch_size=32, shuffle=False,num_workers=4)
  
  def train_one_epoch():
  
      print("Number of batches:", len(dataloader_train))
      total_loss = 0
      stat_loss = 0
  
      total_acc = np.zeros(40)
      stat_acc = np.zeros(40)
  
      model.train()
  
      for batch_idx, batch_data in enumerate(tqdm(dataloader_train)):
          batch_data['image'] = batch_data['image'].to(device)
          batch_data['attributes'] = batch_data['attributes'].to(device)
  
          # Forward pass
          optimizer.zero_grad()
          inputs = batch_data['image']
  
          pred = model(inputs)
          pred_labels = torch.where(pred > 0.5, 1.0, 0.0)
          real_labels = torch.index_select(batch_data['attributes'], 1, torch.tensor(list(range(40))).to(device))
  
          # Compute loss and gradients
          loss = criterion(pred,real_labels)
          acc = compute_accuracy(pred_labels,real_labels)
  
          stat_loss += loss.item()
          total_loss += loss.item()
          stat_acc += acc
          total_acc += acc
  
          loss.backward()
          optimizer.step()
  
          batch_interval = 50
          if (batch_idx+1) % batch_interval == 0:
              log_string(' ---- batch: %03d ----' % (batch_idx+1))
              log_string('mean loss on the last 50 batches: %f'%(stat_loss/batch_interval))
              log_string('mean accuracy on the last 50 batches: '+ str(stat_acc/batch_interval))
              log_string('mean acc. across the different attributes: ' + str(np.mean(stat_acc/batch_interval)))
              stat_loss = 0
              stat_acc = 0
  
      total_mean_loss = total_loss / len(dataloader_train)
      total_mean_acc = total_acc / len(dataloader_train)
      log_string('mean loss over training set: %f' % (total_mean_loss))
      log_string('mean accuracy over training set: ' + str(total_mean_acc))
      log_string('mean acc. across the different attributes: '+str(np.mean(total_mean_acc)))
  
      return total_mean_loss
  
  def evaluate_one_epoch():
  
      model.eval()
  
      total_loss = 0
      stat_loss = 0
      total_acc = np.zeros(40)
      stat_acc = np.zeros(40)
  
      print("Number of batches:", len(dataloader_val))
  
      for batch_idx, batch_data in enumerate(tqdm(dataloader_val)):
          batch_data['image'] = batch_data['image'].to(device)
          batch_data['attributes'] = batch_data['attributes'].to(device)
  
          # Forward pass
          inputs = batch_data['image']
          with torch.no_grad():
              pred = model(inputs)
              pred_labels = torch.where(pred > 0.5 , 1.0,0.0)
  
              real_labels = torch.index_select(batch_data['attributes'], 1, torch.tensor(list(range(40))).to(device))
  
          # Compute loss and metrics
          loss = criterion(pred, real_labels)
          acc = compute_accuracy(pred_labels, real_labels)
  
          stat_loss += loss.item()
          total_loss += loss.item()
          stat_acc += acc
          total_acc += acc
  
          batch_interval = 50
          if (batch_idx + 1) % batch_interval == 0:
              log_string(' ---- batch: %03d ----' % (batch_idx + 1))
              log_string('mean loss on the last 50 batches: %f' % (stat_loss/batch_interval))
              log_string('mean accuracy on the last 50 batches: ' + str(stat_acc/batch_interval))
              log_string('mean acc. across the different attributes: ' + str(np.mean(stat_acc/batch_interval)))
              stat_loss = 0
              stat_acc = 0
  
      total_mean_loss = total_loss / len(dataloader_val)
      total_mean_acc = total_acc / len(dataloader_val)
  
      log_string('mean loss over validation set: %f' % (total_mean_loss))
      log_string('mean accuracy over validation set: ' + str(total_mean_acc))
      log_string('mean acc. across the different attributes: ' + str(np.mean(total_mean_acc)))
  
      return total_mean_loss
  
  # Load model
  model = OracleResnet(weights_path=opt.oracle_pretraining_path, freeze_layers=False, unfreeze_last_block=True)
  model.to(device)
  
  # Prepare optimizer
  if opt.optimizer == 'adam':
      optimizer = torch.optim.Adam(model.parameters(),lr=opt.lr)
  else:
      optimizer = torch.optim.SGD(model.parameters(),lr=opt.lr)
  
  # Prepare loss
  criterion = nn.BCELoss(reduction='mean')
  
  def compute_accuracy(pred, target):
      same_ids = (pred == target).float().cpu()
      return torch.mean(same_ids,axis=0).numpy()
  
  scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=opt.step_size, gamma=opt.gamma_scheduler, verbose=True)
  
  LOG_DIR = os.path.join(config.checkpoint_dir, "oracle_attribute", opt.oracle_name)
  
  if not os.path.exists(LOG_DIR):
      os.mkdir(LOG_DIR)
  
  LOG_FOUT = open(os.path.join(LOG_DIR, 'log_train.txt'), 'a')
  LOG_FOUT.write(str(opt)+'\n')
  
  def log_string(out_str):
      LOG_FOUT.write(out_str+'\n')
      LOG_FOUT.flush()
      print(out_str)
  
  lowest_loss = 100000
  start_epoch=0
  log_string("Starting training from the beginning.")
  
  for epoch in range(start_epoch, opt.num_epochs):
  
      log_string(' **** EPOCH: %03d ****' % (epoch+1))
  
      train_one_epoch()
  
      # Evaluate
      log_string(' **** EVALUATION AFTER EPOCH %03d ****' % (epoch+1))
      total_mean_loss = evaluate_one_epoch()
      if total_mean_loss < lowest_loss:
          lowest_loss = total_mean_loss
          save_dict = {'epoch': epoch+1, 'optimizer_state_dict': optimizer.state_dict(), 'loss': total_mean_loss, 'model_state_dict': model.state_dict()}
          torch.save(save_dict, os.path.join(config.checkpoint_dir, "oracle_attribute", opt.oracle_name, 'checkpoint.tar'))
  
      scheduler.step()
  
  