def train(train_loader, model, ce, optimizer, epoch):
    model.train()
    lambda_ = ((2.0 / (1 + np.exp((((- 10.0) * epoch) / args.epochs)))) - 1.0)
    model.grl.set_lambda(lambda_)
    correct_a = 0.0
    correct_b = 0.0
    total_datasize = 0.0
    total_loss_a = 0.0
    total_loss_b = 0.0
    total_loss_c = 0.0
    total_loss = 0.0
    (ce_criterion, xe_criterion) = ce
    pbar = tqdm(enumerate(train_loader))
    output_softmax = nn.Softmax(dim=1)
    for (batch_idx, (data, label_a, label_b)) in pbar:
        if args.cuda:
            data = data.cuda()
        (data, label_a) = (Variable(data), Variable(label_a))
        label_b = Variable(label_b)
        (logits_spk, feat_spk, logits_dom, feat_dom) = model(data)
        true_labels_a = label_a.cuda()
        true_labels_b = label_b.cuda()
        spk_label = logits_spk
        dom_lable = logits_dom
        if (args.loss_type == 'soft'):
            spk_loss = ce_criterion(logits_spk, true_labels_a)
        elif (args.loss_type == 'asoft'):
            (spk_label, _) = spk_label
            spk_loss = xe_criterion(logits_spk, true_labels_a)
        elif (args.loss_type == 'center'):
            loss_cent = ce_criterion(logits_spk, true_labels_a)
            loss_xent = xe_criterion(feat_spk, true_labels_a)
            spk_loss = ((args.loss_ratio * loss_xent) + loss_cent)
        elif (args.loss_type == 'amsoft'):
            spk_loss = xe_criterion(logits_spk, true_labels_a)
        dom_loss = (args.dom_ratio * ce_criterion(dom_lable, true_labels_b))
        loss = (spk_loss + dom_loss)
        if args.sim_ratio:
            spk_dom_sim_loss = torch.cosine_similarity(feat_spk, feat_dom, dim=1).pow(2).mean()
            spk_dom_sim_loss = (args.sim_ratio * spk_dom_sim_loss)
            loss += spk_dom_sim_loss
        predicted_labels_a = output_softmax(spk_label)
        predicted_one_labels_a = torch.max(predicted_labels_a, dim=1)[1]
        minibatch_correct_a = float((predicted_one_labels_a.cuda() == true_labels_a.cuda()).sum().item())
        minibatch_acc_a = (minibatch_correct_a / len(predicted_one_labels_a))
        correct_a += minibatch_correct_a
        predicted_labels_b = output_softmax(dom_lable)
        predicted_one_labels_b = torch.max(predicted_labels_b, dim=1)[1]
        minibatch_correct_b = float((predicted_one_labels_b.cuda() == true_labels_b.cuda()).sum().item())
        minibatch_acc_b = (minibatch_correct_b / len(predicted_one_labels_b))
        correct_b += minibatch_correct_b
        total_datasize += len(predicted_one_labels_a)
        total_loss_a += float(spk_loss.item())
        total_loss_b += float(dom_loss.item())
        total_loss_c += (float(spk_dom_sim_loss.item()) if args.sim_ratio else 0.0)
        total_loss += float(loss.item())
        optimizer.zero_grad()
        loss.backward()
        if ((args.loss_type == 'center') and (args.loss_ratio != 0)):
            for param in xe_criterion.parameters():
                param.grad.data *= (1.0 / args.loss_ratio)
        optimizer.step()
        if ((batch_idx % args.log_interval) == 0):
            pbar.set_description('Train Epoch {:2d}: [{:4d}/{:4d}({:3.0f}%)] AvgLoss: {:.4f} SpkLoss: {:.4f} DomLoss: {:.4f} SimLoss: {:.4f} Batch Accuracy: Spk: {:.4f}%, Dom: {:.4f}%'.format(epoch, batch_idx, len(train_loader), ((100.0 * batch_idx) / len(train_loader)), (total_loss / (batch_idx + 1)), (total_loss_a / (batch_idx + 1)), (total_loss_b / (batch_idx + 1)), (total_loss_c / (batch_idx + 1)), (100.0 * minibatch_acc_a), (100.0 * minibatch_acc_b)))
    print('\n\x1b[91mTrain Epoch {}: Avg loss: {:.4f} Spk Loss: {:.4f} Dom Loss: {:.4f} .'.format(epoch, (total_loss / len(train_loader)), (total_loss_a / len(train_loader)), (total_loss_b / len(train_loader))))
    print('Spk Accuracy:{:.4f}%, Dom Accuracy:{:.4f}%.\x1b[0m'.format(((100 * correct_a) / total_datasize), ((100 * correct_b) / total_datasize)))
    writer.add_scalar('Train/Spk_Accuracy', (correct_a / total_datasize), epoch)
    writer.add_scalar('Train/Dom_Accuracy', (correct_b / total_datasize), epoch)
    writer.add_scalar('Train/Loss', (total_loss / len(train_loader)), epoch)
    torch.cuda.empty_cache()
def train(train_loader, model, ce, optimizer, epoch):
    model.train()
    lambda_ = ((2.0 / (1 + np.exp((((- 10.0) * epoch) / args.epochs)))) - 1.0)
    model.grl.set_lambda(lambda_)
    correct_a = 0.0
    correct_b = 0.0
    total_datasize = 0.0
    total_loss_a = 0.0
    total_loss_b = 0.0
    total_loss_c = 0.0
    total_loss = 0.0
    (ce_criterion, xe_criterion) = ce
    pbar = tqdm(enumerate(train_loader))
    output_softmax = nn.Softmax(dim=1)
    for (batch_idx, (data, label_a, label_b)) in pbar:
        if args.cuda:
            data = data.cuda()
        (data, label_a) = (Variable(data), Variable(label_a))
        label_b = Variable(label_b)
        (logits_spk, feat_spk, logits_dom, feat_dom) = model(data)
        true_labels_a = label_a.cuda()
        true_labels_b = label_b.cuda()
        spk_label = logits_spk
        dom_lable = logits_dom
        if (args.loss_type == 'soft'):
            spk_loss = ce_criterion(logits_spk, true_labels_a)
        elif (args.loss_type == 'asoft'):
            (spk_label, _) = spk_label
            spk_loss = xe_criterion(logits_spk, true_labels_a)
        elif (args.loss_type == 'center'):
            loss_cent = ce_criterion(logits_spk, true_labels_a)
            loss_xent = xe_criterion(feat_spk, true_labels_a)
            spk_loss = ((args.loss_ratio * loss_xent) + loss_cent)
        elif (args.loss_type == 'amsoft'):
            spk_loss = xe_criterion(logits_spk, true_labels_a)
        dom_loss = (args.dom_ratio * ce_criterion(dom_lable, true_labels_b))
        loss = (spk_loss + dom_loss)
        if args.sim_ratio:
            spk_dom_sim_loss = torch.cosine_similarity(feat_spk, feat_dom, dim=1).pow(2).mean()
            spk_dom_sim_loss = (args.sim_ratio * spk_dom_sim_loss)
            loss += spk_dom_sim_loss
        predicted_labels_a = output_softmax(spk_label)
        predicted_one_labels_a = torch.max(predicted_labels_a, dim=1)[1]
        minibatch_correct_a = float((predicted_one_labels_a.cuda() == true_labels_a.cuda()).sum().item())
        minibatch_acc_a = (minibatch_correct_a / len(predicted_one_labels_a))
        correct_a += minibatch_correct_a
        predicted_labels_b = output_softmax(dom_lable)
        predicted_one_labels_b = torch.max(predicted_labels_b, dim=1)[1]
        minibatch_correct_b = float((predicted_one_labels_b.cuda() == true_labels_b.cuda()).sum().item())
        minibatch_acc_b = (minibatch_correct_b / len(predicted_one_labels_b))
        correct_b += minibatch_correct_b
        total_datasize += len(predicted_one_labels_a)
        total_loss_a += float(spk_loss.item())
        total_loss_b += float(dom_loss.item())
        total_loss_c += (float(spk_dom_sim_loss.item()) if args.sim_ratio else 0.0)
        total_loss += float(loss.item())
        optimizer.zero_grad()
        loss.backward()
        if ((args.loss_type == 'center') and (args.loss_ratio != 0)):
            for param in xe_criterion.parameters():
                param.grad.data *= (1.0 / args.loss_ratio)
        optimizer.step()
        if ((batch_idx % args.log_interval) == 0):
            pbar.set_description('Train Epoch {:2d}: [{:4d}/{:4d}({:3.0f}%)] AvgLoss: {:.4f} SpkLoss: {:.4f} DomLoss: {:.4f} SimLoss: {:.4f} Batch Accuracy: Spk: {:.4f}%, Dom: {:.4f}%'.format(epoch, batch_idx, len(train_loader), ((100.0 * batch_idx) / len(train_loader)), (total_loss / (batch_idx + 1)), (total_loss_a / (batch_idx + 1)), (total_loss_b / (batch_idx + 1)), (total_loss_c / (batch_idx + 1)), (100.0 * minibatch_acc_a), (100.0 * minibatch_acc_b)))
    print('\n\x1b[91mTrain Epoch {}: Avg loss: {:.4f} Spk Loss: {:.4f} Dom Loss: {:.4f} .'.format(epoch, (total_loss / len(train_loader)), (total_loss_a / len(train_loader)), (total_loss_b / len(train_loader))))
    print('Spk Accuracy:{:.4f}%, Dom Accuracy:{:.4f}%.\x1b[0m'.format(((100 * correct_a) / total_datasize), ((100 * correct_b) / total_datasize)))
    writer.add_scalar('Train/Spk_Accuracy', (correct_a / total_datasize), epoch)
    writer.add_scalar('Train/Dom_Accuracy', (correct_b / total_datasize), epoch)
    writer.add_scalar('Train/Loss', (total_loss / len(train_loader)), epoch)
    torch.cuda.empty_cache()