//  [python]
def main():
    opt.seed = random.randint(1, 10000)
    print("Random Seed: ", opt.seed)
    torch.manual_seed(opt.seed)

    torch.cuda.manual_seed_all(opt.seed)
    # if opt.use_GAN:
    #     trans = transforms.Compose(transforms = [   
    #         transforms.ToTensor(),
    #         transforms.Normalize(mean = (0.5,0.5,0.5), std = (0.5,0.5,0.5))
    #     ])
    # else:
    #      trans = transforms.Compose(transforms = [   
    #         transforms.ToTensor()
    #     ])

    print("===> Loading datasets")
    data_path1 = "./data/train_xx1.h5"
    data_path2 = "./data/train_xx2.h5"
    data_path3 = "./data/train_xx3.h5"
    train_set  = DatasetFromHdf5_3_data(data_root_1= data_path1,data_root_2=data_path2,data_root_3=data_path3)

    training_data_loader = DataLoader(dataset=train_set, num_workers=1, batch_size=opt.batch_size, shuffle=True,drop_last=True)
    print("===> Building model")
    wn=None
    if opt.wn==True:
        print('===>use weightNorm')
        wn = lambda x: torch.nn.utils.parametrizations.weight_norm(x)
        model_DB = DeBlockNet(wn=wn, num_sliceBlock=6,is_train=True,input_size=128,use_LSTM=True,use_mask =True,mask_num_sliceBlock=2)
    else:
        model_DB = DeBlockNet(wn=None, num_sliceBlock=6,is_train=True,input_size=128,use_LSTM=True,use_mask =True,mask_num_sliceBlock=2).cuda()
    if opt.use_Warp:
        if opt.flow_type=='flownet2.0':
            print('===>use optical flow(flownet2.0) to warp')                
            model_MC = BiFlowNet.FlowNet2(opt, requires_grad=False)                   
            MC_model_path =opt.pre_model_MC_path
            print("===> Load %s" %MC_model_path)
            checkpoint = torch.load(MC_model_path,map_location=lambda storage, loc: storage)
            model_MC.load_state_dict(checkpoint['state_dict'])
        elif opt.flow_type=='pwc_net':
            print('===>use optical flow(pwc_net) to warp') 
            model_MC = BiFlowNet.PWCNet(requires_grad=False,is_train=False)
            MC_model_path = opt.pre_model_PWC_path
            print("===> Load %s" %MC_model_path)
            data = torch.load(MC_model_path,map_location=lambda storage, loc: storage)
            if 'state_dict' in data.keys():
                model_MC.load_state_dict(data['state_dict'])
            else:
                model_MC.load_state_dict(data)
            model_MC.eval()

    model_DB_path = './pretrained_models/pretrained_model_LD/pretrained_model_LD32_GAN.pth'
    model_DB.load_state_dict(torch.load(model_DB_path,map_location=lambda storage, loc: storage))
    print("===> Load %s" % model_DB_path)

    #G_set_up
    model_MPSN = MPSN(args=opt,motionCompensator=model_MC,U_net=model_DB,use_Warp=opt.use_Warp)
    criterion = nn.MSELoss()
    criterion_L1 = nn.L1Loss()
    print("===> Setting GPU")
    GPU_list = list(map(int,opt.GPUs))
    if len(GPU_list) > 1:
        model_MPSN        = nn.DataParallel(model_MPSN,device_ids=GPU_list)
        model_MPSN        = model_MPSN.cuda()
        criterion = criterion.cuda()
        criterion_L1=criterion_L1.cuda()
    else:
        model_MPSN = model_MPSN.cuda()
        criterion = criterion.cuda()
        criterion_L1=criterion_L1.cuda()
    print("===> Setting Optimizer")
    if opt.solver == 'SGD':
        param_2 = {'params':model_DB.parameters(),'lr':opt.lr_init}
        param_groups = [param_2]
        optimizer = optim.SGD(param_groups)
    elif opt.solver == 'ADAM':
        param_2 = {'params':model_DB.parameters(),'lr':opt.lr_init}
        param_groups = [param_2]
        optimizer = optim.Adam(param_groups)
    else:
        raise Exception("Not supported solver (%s)" %opt.solver)

    #D_set_up
    model_D = None
    if opt.use_GAN:
        print('===> use GAN')
        model_D = define_D(opt)
        if len(GPU_list) > 1:
            model_D = nn.DataParallel(model_D,device_ids=GPU_list).cuda()
        else:
            model_D=model_D.cuda()
            
        if opt.solver == 'SGD':
            param_D = {'params':model_D.parameters(),'lr':opt.D_lr_init}   
            optimizer_D = optim.SGD([param_D])
        elif opt.solver == 'ADAM':
            param_D = {'params':model_D.parameters(),'lr':opt.D_lr_init}
            optimizer_D = optim.Adam([param_D])
        else:
            raise Exception("Not supported solver (%s)" %opt.solver)

    print("===>start  Training")
    for epoch in range(opt.start_epoch, opt.nEpochs + 1):
        lr,lr_D = adjust_learning_rate(epoch-1)
        optimizer.param_groups[0]['lr']  = lr
        optimizer_D.param_groups[0]['lr'] = lr_D
        if opt.use_GAN:
            print("Epoch = {}, lr = {} lr_D={}".format(epoch, optimizer.param_groups[0]["lr"],optimizer_D.param_groups[0]["lr"]))
        else:
            print("Epoch = {}, lr = {} ".format(epoch, optimizer.param_groups[0]["lr"]))
        total_count=0
        if opt.use_GAN:
            total_count=train_GAN(training_data_loader, optimizer, model_MPSN,optimizer_D, model_D, criterion,criterion_L1, epoch,total_count)
        else:
            # total_count=train(training_data_loader, optimizer, model_MPSN, criterion,criterion_L1, epoch,total_count)
            raise Exception("Not supported GAN" )
        save_checkpoint(model_D , model_DB,epoch)