""" Barlow Twin self-supervision training
"""
import argparse
import json
import math
import os
import random
import signal
import subprocess
import sys
import time
from tqdm import tqdm
from torch import nn, optim
import torch
import torchvision
import torchinfo
sys.path.append(os.getcwd())
import utilities.runUtils as rutl
import utilities.logUtils as lutl
from algorithms.autoencoder import AutoEncoder
from datacode.natural_image_data import Cifar100Dataset
from datacode.ultrasound_data import FetalUSFramesDataset
from datacode.augmentations import AEncStandardTransform, AEncInpaintTransform
print(f"Pytorch version: {torch.__version__}")
print(f"cuda version: {torch.version.cuda}")
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print("Device Used:", device)
###============================= Configure and Setup ===========================
CFG = rutl.ObjDict(
use_amp = True, #automatic Mixed precision
datapath = "/home/USR/WERK/data/",
valdatapath = "/home/USR/WERK/valdata/",
skip_count = 5,
epochs = 1000,
batch_size = 2048,
workers = 16,
image_size = 256,
learning_rate = 1e-3,
weight_decay = 1e-6,
sched_step = 50, ## epoch
sched_gamma = 0.5624, # 1/10 every 200
autoenc_map = "standard", # standard, denoise, inpaint
featx_arch = "resnet50", # "resnet34/50/101"
featx_pretrain = "IMAGENET-1K", # "IMAGENET-1K" or None
print_freq_step = 1000, #steps
ckpt_freq_epoch = 5, #epochs
valid_freq_epoch = 5, #epochs
disable_tqdm = False, #True--> to disable
checkpoint_dir = "hypotheses/-dummy/ssl-autoenc",
resume_training = False,
)
## --------
parser = argparse.ArgumentParser(description='Auto Encoder architecture training Training')
parser.add_argument('--load-json', type=str, metavar='JSON',
help='Load settings from file in json format. Command line options override values in file.')
args = parser.parse_args()
if args.load_json:
with open(args.load_json, 'rt') as f:
CFG.__dict__.update(json.load(f))
### ----------------------------------------------------------------------------
CFG.gLogPath = CFG.checkpoint_dir
CFG.gWeightPath = CFG.checkpoint_dir + '/weights/'
### ============================================================================
def getDataLoaders():
if CFG.autoenc_map == "standard":
transform_obj = AEncStandardTransform(image_size=CFG.image_size)
elif CFG.autoenc_map == "inpaint":
transform_obj = AEncInpaintTransform(image_size=CFG.image_size)
else:
raise Exception("Unknown Auto Encoder augmentatoion")
traindataset = FetalUSFramesDataset( hdf5_file= CFG.datapath,
transform = transform_obj,
load2ram = False, frame_skip=CFG.skip_count)
trainloader = torch.utils.data.DataLoader( traindataset, shuffle=True,
batch_size=CFG.batch_size, num_workers=CFG.workers,
pin_memory=True)
validdataset = FetalUSFramesDataset( hdf5_file= CFG.valdatapath,
transform = transform_obj,
load2ram = False, frame_skip=CFG.skip_count)
validloader = torch.utils.data.DataLoader( validdataset, shuffle=False,
batch_size=CFG.batch_size, num_workers=CFG.workers,
pin_memory=True)
lutl.LOG2DICTXT({"TRAIN DatasetClass":traindataset.get_info(),
"TransformsClass": str(transform_obj.get_composition()),
}, CFG.gLogPath +'/misc.txt')
lutl.LOG2DICTXT({"VALID DatasetClass":validdataset.get_info(),
"TransformsClass": str(transform_obj.get_composition()),
}, CFG.gLogPath +'/misc.txt')
return trainloader, validloader
def getModelnOptimizer():
model = AutoEncoder(arch=CFG.featx_arch,
pretrained=CFG.featx_pretrain).to(device)
optimizer = optim.AdamW(model.parameters(), lr=CFG.learning_rate,
weight_decay=CFG.weight_decay)
scheduler = optim.lr_scheduler.StepLR(optimizer,
step_size=CFG.sched_step, gamma=CFG.sched_gamma)
model_info = torchinfo.summary(model, (1, 3, CFG.image_size, CFG.image_size),
verbose=0)
lutl.LOG2TXT(model_info, CFG.gLogPath +'/misc.txt', console= False)
return model.to(device), optimizer, scheduler
def getLossFunc():
mse = nn.MSELoss()
# def scaledMSE(pred, tgt):
# loss = mse(pred, tgt) *256
# return loss
return mse
def simple_main():
### SETUP
rutl.START_SEED()
torch.cuda.device(device)
torch.backends.cudnn.benchmark = True
if os.path.exists(CFG.checkpoint_dir) and (not CFG.resume_training):
raise Exception("CheckPoint folder already exists and restart_training not enabled; Somethings Wrong!")
if not os.path.exists(CFG.gWeightPath): os.makedirs(CFG.gWeightPath)
with open(CFG.gLogPath+"/exp_config.json", 'a') as f:
json.dump(vars(CFG), f, indent=4)
### DATA ACCESS
trainloader, validloader = getDataLoaders()
### MODEL, OPTIM
model, optimizer, scheduler = getModelnOptimizer()
lossfn = getLossFunc()
## Automatically resume from checkpoint if it exists and enabled
ckpt = None
if CFG.resume_training:
try: ckpt = torch.load(CFG.gWeightPath+'/checkpoint-1.pth', map_location='cpu')
except:
try:ckpt = torch.load(CFG.gWeightPath+'/checkpoint-0.pth', map_location='cpu')
except: print("Check points are not loadable. Starting fresh...")
if ckpt:
start_epoch = ckpt['epoch']
model.load_state_dict(ckpt['model'])
optimizer.load_state_dict(ckpt['optimizer'])
lutl.LOG2TXT(f"Restarting Training from EPOCH:{start_epoch} of {CFG.checkpoint_dir}", CFG.gLogPath +'/misc.txt')
else:
start_epoch = 0
### MODEL TRAINING
start_time = time.time()
best_loss = float('inf')
wgt_suf = 0 # foolproof savetime crash
if CFG.use_amp: scaler = torch.cuda.amp.GradScaler() # for mixed precision
for epoch in range(start_epoch, CFG.epochs):
## ---- Training Routine ----
t_running_loss_ = 0
model.train()
for step, (y1, y2) in tqdm(enumerate(trainloader,
start=epoch * len(trainloader)),
disable=CFG.disable_tqdm):
y1 = y1.to(device, non_blocking=True)
y2 = y2.to(device, non_blocking=True)
optimizer.zero_grad()
if CFG.use_amp: ## with mixed precision
with torch.cuda.amp.autocast():
y_pred = model.forward(y1)
loss = lossfn(y_pred, y2)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
else:
y_pred = model.forward(y1)
loss = lossfn(y_pred, y2)
loss.backward()
optimizer.step()
t_running_loss_+=loss.item()
if step % CFG.print_freq_step == 0:
stats = dict(epoch=epoch, step=step,
step_loss=loss.item(),
time=int(time.time() - start_time))
lutl.LOG2DICTXT(stats, CFG.checkpoint_dir +'/train-stats.txt')
train_epoch_loss = t_running_loss_/len(trainloader)
if scheduler: scheduler.step()
# save checkpoint
if (epoch+1) % CFG.ckpt_freq_epoch == 0:
wgt_suf = (wgt_suf+1) %2
state = dict(epoch=epoch, model=model.state_dict(),
optimizer=optimizer.state_dict())
torch.save(state, CFG.gWeightPath +f'/checkpoint-{wgt_suf}.pth')
## ---- Validation Routine ----
if (epoch+1) % CFG.valid_freq_epoch == 0:
model.eval()
v_running_loss_ = 0
with torch.no_grad():
for (y1, y2) in tqdm(validloader, total=len(validloader),
disable=CFG.disable_tqdm):
y1 = y1.to(device, non_blocking=True)
y2 = y2.to(device, non_blocking=True)
y_pred = model.forward(y1)
loss = lossfn(y_pred, y2)
v_running_loss_ += loss.item()
valid_epoch_loss = v_running_loss_/len(validloader)
best_flag = False
if valid_epoch_loss < best_loss:
best_flag = True
best_loss = valid_epoch_loss
v_stats = dict(epoch=epoch, best=best_flag, wgt_suf=wgt_suf,
train_loss=train_epoch_loss,
valid_loss=valid_epoch_loss)
lutl.LOG2DICTXT(v_stats, CFG.gLogPath+'/valid-stats.txt')
if __name__ == '__main__':
simple_main()
