def main(args):

    def log_string(str):
        logger.info(str)
        print(str)
    'HYPER PARAMETER'
    os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
    'CREATE DIR'
    timestr = str(datetime.datetime.now().strftime('%Y-%m-%d_%H-%M'))
    exp_dir = Path('./log/')
    exp_dir.mkdir(exist_ok=True)
    exp_dir = exp_dir.joinpath('classification')
    exp_dir.mkdir(exist_ok=True)
    if (args.log_dir is None):
        exp_dir = exp_dir.joinpath(timestr)
    else:
        exp_dir = exp_dir.joinpath(args.log_dir)
    exp_dir.mkdir(exist_ok=True)
    checkpoints_dir = exp_dir.joinpath('checkpoints/')
    checkpoints_dir.mkdir(exist_ok=True)
    log_dir = exp_dir.joinpath('logs/')
    log_dir.mkdir(exist_ok=True)
    'LOG'
    args = parse_args()
    logger = logging.getLogger('Model')
    logger.setLevel(logging.INFO)
    formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    file_handler = logging.FileHandler(('%s/%s.txt' % (log_dir, args.model)))
    file_handler.setLevel(logging.INFO)
    file_handler.setFormatter(formatter)
    logger.addHandler(file_handler)
    log_string('PARAMETER ...')
    log_string(args)
    'DATA LOADING'
    log_string('Load dataset ...')
    visual_data_path = 'data/visual_data_pcd/'
    tactile_data_path = 'data/tactile_pcd_10_sampled_21.02/'
    train_dataset = PCDPointCloudData(visual_data_path, folder='Train', sample_method='Voxel', num_point=args.num_point, sample=True, rotation=False, est_normal=args.use_normals)
    test_dataset = PCDPointCloudData(visual_data_path, folder='Test', sample_method='Voxel', num_point=args.num_point, sample=True, rotation=False, est_normal=args.use_normals)
    if (args.random_choose_sparse is True):
        raise NotImplementedError('Function Not Implemented')
    else:
        domain_adaptation_dataset = PCDPointCloudData(tactile_data_path, folder='Train', sample_method='Voxel', num_point=args.num_sparse_point, sample=True, rotation=False, est_normal=args.use_normals)
    trainDataLoader = torch.utils.data.DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=10, drop_last=True)
    domainAdaptationDataLoader = torch.utils.data.DataLoader(domain_adaptation_dataset, batch_size=args.batch_size, shuffle=True, num_workers=10, drop_last=True)
    testDataLoader = torch.utils.data.DataLoader(test_dataset, batch_size=args.batch_size, shuffle=False, num_workers=10)
    'Output middle layers'
    activation = {}

    def get_activation(name):

        def hook(model, input, output):
            activation[name] = output.detach()
        return hook
    'MODEL LOADING'
    num_class = args.num_category
    model = importlib.import_module(args.model)
    shutil.copy(('./models/%s.py' % args.model), str(exp_dir))
    shutil.copy('models/pointnet_cls.py', str(exp_dir))
    shutil.copy('data_utils/PCDLoader.py', str(exp_dir))
    shutil.copy('./train_realMulti-DA-Loss_classification.py', str(exp_dir))
    classifier = model.get_model(num_class, normal_channel=args.use_normals)
    criterion = model.get_loss()
    if (args.DA_method == 'coral'):
        criterion_DA = model.get_coral_loss(DA_alpha=args.alpha, DA_lamda=args.lamda)
    elif (args.DA_method == 'mmd'):
        criterion_DA = model.get_mmd_loss(DA_alpha=args.alpha, DA_lamda=args.lamda)
    elif (args.DA_method == 'coral_mmd'):
        criterion_DA = model.get_coral_mmd_loss(DA_alpha=args.alpha, DA_beta=args.beta, DA_lamda=args.lamda)
    elif (args.DA_method == 'multi_coral_mmd'):
        criterion_DA = model.get_multiLayer_loss(DA_alpha=args.alpha, DA_beta=args.beta, DA_lamda=args.lamda)
    else:
        raise NameError('Wrong input for DA method name!')
    classifier.apply(inplace_relu)
    if (not args.use_cpu):
        classifier = classifier.cuda()
        criterion = criterion.cuda()
        criterion_DA = criterion_DA.cuda()
    try:
        checkpoint = torch.load((str(exp_dir) + '/checkpoints/best_model.pth'))
        start_epoch = checkpoint['epoch']
        classifier.load_state_dict(checkpoint['model_state_dict'])
        log_string('Use pretrain model')
    except:
        log_string('No existing model, starting training from scratch...')
        start_epoch = 0
    try:
        min_loss = checkpoint['loss']
        log_string(('Loading model with DA loss %f' % min_loss))
    except:
        log_string('No DA loss found in the model')
        min_loss = 10000.0
    if (args.optimizer == 'Adam'):
        optimizer = torch.optim.Adam(classifier.parameters(), lr=args.learning_rate, betas=(0.9, 0.999), eps=1e-08, weight_decay=args.decay_rate)
    else:
        optimizer = torch.optim.SGD(classifier.parameters(), lr=0.01, momentum=0.9)
    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=20, gamma=0.7)
    global_epoch = 0
    global_step = 0
    best_instance_acc = 0.0
    best_class_acc = 0.0
    running_loss = 0.0
    'TRANING'
    logger.info('Start training...')
    end_epoch = (start_epoch + args.epoch)
    print('start epoch: ', start_epoch)
    print('end epoch: ', end_epoch)
    for epoch in range(start_epoch, end_epoch):
        log_string(('Epoch %d (%d/%s):' % ((global_epoch + 1), (epoch + 1), end_epoch)))
        mean_correct = []
        for (name, param) in classifier.named_parameters():
            if ('feat' in name):
                param.requires_grad = False
        scheduler.step()
        for (batch_id, (data, data_DA)) in tqdm(enumerate(zip(trainDataLoader, domainAdaptationDataLoader), 0), total=len(trainDataLoader), smoothing=0.9):
            optimizer.zero_grad()
            (points, target) = (data['pointcloud'].to(device).float(), data['category'].to(device))
            points_DA = data_DA['pointcloud'].to(device).float()
            points = points.data.cpu().numpy()
            points = provider.random_point_dropout(points)
            points[:, :, 0:3] = provider.random_scale_point_cloud(points[:, :, 0:3])
            points[:, :, 0:3] = provider.shift_point_cloud(points[:, :, 0:3])
            points = torch.Tensor(points)
            points = points.transpose(2, 1)
            points_DA = points_DA.data.cpu().numpy()
            points_DA = provider.random_point_dropout(points_DA)
            points_DA[:, :, 0:3] = provider.random_scale_point_cloud(points_DA[:, :, 0:3])
            points_DA[:, :, 0:3] = provider.shift_point_cloud(points_DA[:, :, 0:3])
            points_DA = torch.Tensor(points_DA)
            points_DA = points_DA.transpose(2, 1)
            if (not args.use_cpu):
                (points, target) = (points.cuda(), target.cuda())
                points_DA = points_DA.cuda()
            (pred, trans_feat) = classifier(points)
            classifier.fc1.register_forward_hook(get_activation('fc1'))
            output_dense_1 = classifier(points)
            feature_dense_1 = activation['fc1']
            classifier.fc1.register_forward_hook(get_activation('fc1'))
            output_DA_1 = classifier(points_DA)
            feature_DA_1 = activation['fc1']
            classifier.fc2.register_forward_hook(get_activation('fc2'))
            output_dense_2 = classifier(points)
            feature_dense_2 = activation['fc2']
            classifier.fc2.register_forward_hook(get_activation('fc2'))
            output_DA_2 = classifier(points_DA)
            feature_DA_2 = activation['fc2']
            (DA_loss, loss) = criterion_DA(pred, target.long(), trans_feat, feature_dense_1, feature_DA_1, feature_dense_2, feature_DA_2)
            pred_choice = pred.data.max(1)[1]
            correct = pred_choice.eq(target.long().data).cpu().sum()
            mean_correct.append((correct.item() / float(points.size()[0])))
            loss.backward()
            optimizer.step()
            global_step += 1
            running_loss += DA_loss.item()
            if ((batch_id % 100) == 99):
                calculate_loss = (running_loss / 100)
                log_string('Running DA loss {} '.format(calculate_loss))
                if (calculate_loss < min_loss):
                    min_loss = calculate_loss
                    logger.info('Save model...')
                    savepath = (str(checkpoints_dir) + '/best_model.pth')
                    log_string(('Saving at %s' % savepath))
                    state = {'epoch': epoch, 'loss': calculate_loss, 'model_state_dict': classifier.state_dict(), 'optimizer_state_dict': optimizer.state_dict()}
                    torch.save(state, savepath)
                running_loss = 0.0
        train_instance_acc = np.mean(mean_correct)
        log_string(('Train Instance Accuracy: %f' % train_instance_acc))
        with torch.no_grad():
            (instance_acc, class_acc) = test(classifier.eval(), testDataLoader, num_class=num_class)
            if (instance_acc >= best_instance_acc):
                best_instance_acc = instance_acc
                best_epoch = (epoch + 1)
            if (class_acc >= best_class_acc):
                best_class_acc = class_acc
            log_string(('Test Instance Accuracy: %f, Class Accuracy: %f' % (instance_acc, class_acc)))
            log_string(('Best Instance Accuracy: %f, Class Accuracy: %f' % (best_instance_acc, best_class_acc)))
            if (instance_acc >= best_instance_acc):
                logger.info('Model will not be saved with vision validation')
            global_epoch += 1
    logger.info('End of training...')
def main(args):

    def log_string(str):
        logger.info(str)
        print(str)
    'HYPER PARAMETER'
    os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
    'CREATE DIR'
    timestr = str(datetime.datetime.now().strftime('%Y-%m-%d_%H-%M'))
    exp_dir = Path('./log/')
    exp_dir.mkdir(exist_ok=True)
    exp_dir = exp_dir.joinpath('classification')
    exp_dir.mkdir(exist_ok=True)
    if (args.log_dir is None):
        exp_dir = exp_dir.joinpath(timestr)
    else:
        exp_dir = exp_dir.joinpath(args.log_dir)
    exp_dir.mkdir(exist_ok=True)
    checkpoints_dir = exp_dir.joinpath('checkpoints/')
    checkpoints_dir.mkdir(exist_ok=True)
    log_dir = exp_dir.joinpath('logs/')
    log_dir.mkdir(exist_ok=True)
    'LOG'
    args = parse_args()
    logger = logging.getLogger('Model')
    logger.setLevel(logging.INFO)
    formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    file_handler = logging.FileHandler(('%s/%s.txt' % (log_dir, args.model)))
    file_handler.setLevel(logging.INFO)
    file_handler.setFormatter(formatter)
    logger.addHandler(file_handler)
    log_string('PARAMETER ...')
    log_string(args)
    'DATA LOADING'
    log_string('Load dataset ...')
    visual_data_path = 'data/visual_data_pcd/'
    tactile_data_path = 'data/tactile_pcd_10_sampled_21.02/'
    train_dataset = PCDPointCloudData(visual_data_path, folder='Train', sample_method='Voxel', num_point=args.num_point, sample=True, rotation=False, est_normal=args.use_normals)
    test_dataset = PCDPointCloudData(visual_data_path, folder='Test', sample_method='Voxel', num_point=args.num_point, sample=True, rotation=False, est_normal=args.use_normals)
    if (args.random_choose_sparse is True):
        raise NotImplementedError('Function Not Implemented')
    else:
        domain_adaptation_dataset = PCDPointCloudData(tactile_data_path, folder='Train', sample_method='Voxel', num_point=args.num_sparse_point, sample=True, rotation=False, est_normal=args.use_normals)
    trainDataLoader = torch.utils.data.DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=10, drop_last=True)
    domainAdaptationDataLoader = torch.utils.data.DataLoader(domain_adaptation_dataset, batch_size=args.batch_size, shuffle=True, num_workers=10, drop_last=True)
    testDataLoader = torch.utils.data.DataLoader(test_dataset, batch_size=args.batch_size, shuffle=False, num_workers=10)
    'Output middle layers'
    activation = {}

    def get_activation(name):

        def hook(model, input, output):
            activation[name] = output.detach()
        return hook
    'MODEL LOADING'
    num_class = args.num_category
    model = importlib.import_module(args.model)
    shutil.copy(('./models/%s.py' % args.model), str(exp_dir))
    shutil.copy('models/pointnet_cls.py', str(exp_dir))
    shutil.copy('data_utils/PCDLoader.py', str(exp_dir))
    shutil.copy('./train_realMulti-DA-Loss_classification.py', str(exp_dir))
    classifier = model.get_model(num_class, normal_channel=args.use_normals)
    criterion = model.get_loss()
    if (args.DA_method == 'coral'):
        criterion_DA = model.get_coral_loss(DA_alpha=args.alpha, DA_lamda=args.lamda)
    elif (args.DA_method == 'mmd'):
        criterion_DA = model.get_mmd_loss(DA_alpha=args.alpha, DA_lamda=args.lamda)
    elif (args.DA_method == 'coral_mmd'):
        criterion_DA = model.get_coral_mmd_loss(DA_alpha=args.alpha, DA_beta=args.beta, DA_lamda=args.lamda)
    elif (args.DA_method == 'multi_coral_mmd'):
        criterion_DA = model.get_multiLayer_loss(DA_alpha=args.alpha, DA_beta=args.beta, DA_lamda=args.lamda)
    else:
        raise NameError('Wrong input for DA method name!')
    classifier.apply(inplace_relu)
    if (not args.use_cpu):
        classifier = classifier.cuda()
        criterion = criterion.cuda()
        criterion_DA = criterion_DA.cuda()
    try:
        checkpoint = torch.load((str(exp_dir) + '/checkpoints/best_model.pth'))
        start_epoch = checkpoint['epoch']
        classifier.load_state_dict(checkpoint['model_state_dict'])
        log_string('Use pretrain model')
    except:
        log_string('No existing model, starting training from scratch...')
        start_epoch = 0
    try:
        min_loss = checkpoint['loss']
        log_string(('Loading model with DA loss %f' % min_loss))
    except:
        log_string('No DA loss found in the model')
        min_loss = 10000.0
    if (args.optimizer == 'Adam'):
        optimizer = torch.optim.Adam(classifier.parameters(), lr=args.learning_rate, betas=(0.9, 0.999), eps=1e-08, weight_decay=args.decay_rate)
    else:
        optimizer = torch.optim.SGD(classifier.parameters(), lr=0.01, momentum=0.9)
    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=20, gamma=0.7)
    global_epoch = 0
    global_step = 0
    best_instance_acc = 0.0
    best_class_acc = 0.0
    running_loss = 0.0
    'TRANING'
    logger.info('Start training...')
    end_epoch = (start_epoch + args.epoch)
    print('start epoch: ', start_epoch)
    print('end epoch: ', end_epoch)
    for epoch in range(start_epoch, end_epoch):
        log_string(('Epoch %d (%d/%s):' % ((global_epoch + 1), (epoch + 1), end_epoch)))
        mean_correct = []
        for (name, param) in classifier.named_parameters():
            if ('feat' in name):
                param.requires_grad = False
        scheduler.step()
        for (batch_id, (data, data_DA)) in tqdm(enumerate(zip(trainDataLoader, domainAdaptationDataLoader), 0), total=len(trainDataLoader), smoothing=0.9):
            optimizer.zero_grad()
            (points, target) = (data['pointcloud'].to(device).float(), data['category'].to(device))
            points_DA = data_DA['pointcloud'].to(device).float()
            points = points.data.cpu().numpy()
            points = provider.random_point_dropout(points)
            points[:, :, 0:3] = provider.random_scale_point_cloud(points[:, :, 0:3])
            points[:, :, 0:3] = provider.shift_point_cloud(points[:, :, 0:3])
            points = torch.Tensor(points)
            points = points.transpose(2, 1)
            points_DA = points_DA.data.cpu().numpy()
            points_DA = provider.random_point_dropout(points_DA)
            points_DA[:, :, 0:3] = provider.random_scale_point_cloud(points_DA[:, :, 0:3])
            points_DA[:, :, 0:3] = provider.shift_point_cloud(points_DA[:, :, 0:3])
            points_DA = torch.Tensor(points_DA)
            points_DA = points_DA.transpose(2, 1)
            if (not args.use_cpu):
                (points, target) = (points.cuda(), target.cuda())
                points_DA = points_DA.cuda()
            (pred, trans_feat) = classifier(points)
            classifier.fc1.register_forward_hook(get_activation('fc1'))
            output_dense_1 = classifier(points)
            feature_dense_1 = activation['fc1']
            classifier.fc1.register_forward_hook(get_activation('fc1'))
            output_DA_1 = classifier(points_DA)
            feature_DA_1 = activation['fc1']
            classifier.fc2.register_forward_hook(get_activation('fc2'))
            output_dense_2 = classifier(points)
            feature_dense_2 = activation['fc2']
            classifier.fc2.register_forward_hook(get_activation('fc2'))
            output_DA_2 = classifier(points_DA)
            feature_DA_2 = activation['fc2']
            (DA_loss, loss) = criterion_DA(pred, target.long(), trans_feat, feature_dense_1, feature_DA_1, feature_dense_2, feature_DA_2)
            pred_choice = pred.data.max(1)[1]
            correct = pred_choice.eq(target.long().data).cpu().sum()
            mean_correct.append((correct.item() / float(points.size()[0])))
            loss.backward()
            optimizer.step()
            global_step += 1
            running_loss += DA_loss.item()
            if ((batch_id % 100) == 99):
                calculate_loss = (running_loss / 100)
                log_string('Running DA loss {} '.format(calculate_loss))
                if (calculate_loss < min_loss):
                    min_loss = calculate_loss
                    logger.info('Save model...')
                    savepath = (str(checkpoints_dir) + '/best_model.pth')
                    log_string(('Saving at %s' % savepath))
                    state = {'epoch': epoch, 'loss': calculate_loss, 'model_state_dict': classifier.state_dict(), 'optimizer_state_dict': optimizer.state_dict()}
                    torch.save(state, savepath)
                running_loss = 0.0
        train_instance_acc = np.mean(mean_correct)
        log_string(('Train Instance Accuracy: %f' % train_instance_acc))
        with torch.no_grad():
            (instance_acc, class_acc) = test(classifier.eval(), testDataLoader, num_class=num_class)
            if (instance_acc >= best_instance_acc):
                best_instance_acc = instance_acc
                best_epoch = (epoch + 1)
            if (class_acc >= best_class_acc):
                best_class_acc = class_acc
            log_string(('Test Instance Accuracy: %f, Class Accuracy: %f' % (instance_acc, class_acc)))
            log_string(('Best Instance Accuracy: %f, Class Accuracy: %f' % (best_instance_acc, best_class_acc)))
            if (instance_acc >= best_instance_acc):
                logger.info('Model will not be saved with vision validation')
            global_epoch += 1
    logger.info('End of training...')