train.py

import os
import json
import argparse
import torch
import dataloaders
import models
import math
from trainer import Trainer
import torch.nn.functional as F
from utils.losses import  CE_loss, Alg_loss
torch.set_num_threads(4)



def main(config, method, resume, gpu, percent):
    os.environ["CUDA_VISIBLE_DEVICES"] = gpu
    torch.manual_seed(42)

    # method = config['model']['method'] 
    config['model']['method'] = method

    backbone = config['model']['backbone'] 
    config['percent'] = percent 
    config['experim_name'] = config['experim_name'].replace('method', method)
    config['experim_name'] = config['experim_name'].replace('percent', str(config['percent'] ))
    config['trainer']['save_dir'] = config['trainer']['save_dir'].replace('backbone', backbone)
    print (config)

    # DATA LOADERS
    config['train_supervised']['percnt_lbl'] = config["percent"]
    config['train_unsupervised']['percnt_lbl'] = config["percent"]
    print ('Sup Loader: ')
    supervised_loader = dataloaders.CDDataset(config['train_supervised'])
    print ('Unsup Loader: ')
    unsupervised_loader = dataloaders.CDDataset(config['train_unsupervised'])
    print ('supervised: ', len(supervised_loader))
    print ('unsupervised: ', len(unsupervised_loader))
    val_loader = dataloaders.CDDataset(config['val_loader'])
    iter_per_epoch = len(unsupervised_loader)

    # SUPERVISED LOSS

    # MODEL
    if backbone == 'ResNet50':
        model = models.FPA_ResNet50_CD(num_classes=val_loader.dataset.num_classes, 
                                        conf=config['model'],
                                        len_unsper=len(unsupervised_loader))
    print(f'\n{model}\n')

    # TRAINING
    trainer = Trainer(
        model=model,
        resume=resume,
        config=config,
        supervised_loader=supervised_loader,
        unsupervised_loader=unsupervised_loader,
        val_loader=val_loader,
        iter_per_epoch=iter_per_epoch)

    trainer.train()

if __name__=='__main__':
    # PARSE THE ARGS
    parser = argparse.ArgumentParser(description='PyTorch Training')
    parser.add_argument('-c', '--config', default='configs/config_dataset.json',type=str,
                        help='Path to the config file')
    parser.add_argument('-d', '--dataset', default='LEVIR',type=str,
                        help='Path to the config file')    
    parser.add_argument('-m', '--method', default='Base3+RA+RRP',type=str,
                        help='test method')
    parser.add_argument('-r', '--resume', default=None, type=str,
                        help='Path to the .pth model checkpoint to resume training')
    parser.add_argument('-g', '--gpu', default=1, type=int,
                        help='indices of GPUs to enable (default: all)')
    parser.add_argument('-p', '--percent', default=5, type=int,
                        help='percent of the labeled training set')
    args = parser.parse_args()

    torch.backends.cudnn.benchmark = True

    gpu = str(args.gpu)
    method = args.method
    percent = args.percent
    config = json.load(open(args.config.replace('dataset', args.dataset)))
    if args.resume != None:
        args.resume = args.resume.replace('method', args.method)\
                                 .replace('percent', str(args.percent))\
                                 .replace('dataset', args.dataset)
    main(config, method, args.resume, gpu, percent)