Ejemplo n.º 1
0
def train(model, tokenizer, train_data, valid_data, args, eos=False):
    model.train()

    train_dataset = TextDataset(train_data)
    train_dataloader = DataLoader(train_dataset, sampler=RandomSampler(train_dataset),
                                  batch_size=args.train_batch_size, num_workers=args.num_workers,
                                  collate_fn=lambda x: collate_fn(x, tokenizer, args.max_seq_length, eos=eos, tokenizer_type=args.tokenizer))

    valid_dataset = TextDataset(valid_data)
    valid_dataloader = DataLoader(valid_dataset, sampler=SequentialSampler(valid_dataset),
                                  batch_size=args.eval_batch_size, num_workers=args.num_workers,
                                  collate_fn=lambda x: collate_fn(x, tokenizer, args.max_seq_length, eos=eos, tokenizer_type=args.tokenizer))

    valid_noisy = [x['noisy'] for x in valid_data]
    valid_clean = [x['clean'] for x in valid_data]

    epochs = (args.max_steps - 1) // len(train_dataloader) + 1
    optimizer = torch.optim.Adam(model.parameters(), lr=args.lr,
                                 betas=eval(args.adam_betas), eps=args.eps,
                                 weight_decay=args.weight_decay)
    lr_lambda = lambda x: x / args.num_warmup_steps if x <= args.num_warmup_steps else (x / args.num_warmup_steps) ** -0.5
    scheduler = LambdaLR(optimizer, lr_lambda)

    step = 0
    best_val_gleu = -float("inf")
    meter = Meter()
    for epoch in range(1, epochs + 1):
        print("===EPOCH: ", epoch)
        for batch in train_dataloader:
            step += 1
            batch = tuple(t.to(args.device) for t in batch)
            loss, items = calc_loss(model, batch)
            meter.add(*items)

            loss.backward()
            if args.max_grad_norm > 0:
                nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)
            optimizer.step()
            model.zero_grad()
            scheduler.step()

            if step % args.log_interval == 0:
                lr = scheduler.get_lr()[0]
                loss_sent, loss_token = meter.average()

                logger.info(f' [{step:5d}] lr {lr:.6f} | {meter.print_str(True)}')
                nsml.report(step=step, scope=locals(), summary=True,
                            train__lr=lr, train__loss_sent=loss_sent, train__token_ppl=math.exp(loss_token))
                meter.init()

            if step % args.eval_interval == 0:
                start_eval = time.time()
                (val_loss, val_loss_token), valid_str = evaluate(model, valid_dataloader, args)
                prediction = correct(model, tokenizer, valid_noisy, args, eos=eos, length_limit=0.1)
                val_em = em(prediction, valid_clean)
                cnt = 0
                for noisy, pred, clean in zip(valid_noisy, prediction, valid_clean):
                    print(f'[{noisy}], [{pred}], [{clean}]')
                    # 10개만 출력하기
                    cnt += 1
                    if cnt == 20:
                        break
                val_gleu = gleu(prediction, valid_clean)

                logger.info('-' * 89)
                logger.info(f' [{step:6d}] valid | {valid_str} | em {val_em:5.2f} | gleu {val_gleu:5.2f}')
                logger.info('-' * 89)
                nsml.report(step=step, scope=locals(), summary=True,
                            valid__loss_sent=val_loss, valid__token_ppl=math.exp(val_loss_token),
                            valid__em=val_em, valid__gleu=val_gleu)

                if val_gleu > best_val_gleu:
                    best_val_gleu = val_gleu
                    nsml.save("best")
                meter.start += time.time() - start_eval

            if step >= args.max_steps:
                break
        #nsml.save(epoch)
        if step >= args.max_steps:
            break
Ejemplo n.º 2
0
def main():
    global opts
    opts = parser.parse_args()
    opts.cuda = 0

    # Set GPU
    seed = opts.seed
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)

    os.environ['CUDA_VISIBLE_DEVICES'] = opts.gpu_ids
    use_gpu = torch.cuda.is_available()
    if use_gpu:
        opts.cuda = 1
        print("Number of device: {}".format(torch.cuda.device_count()))
        print("Currently using GPU {}".format(opts.gpu_ids))
        cudnn.benchmark = True
        torch.cuda.manual_seed_all(seed)
    else:
        print("Currently using CPU (GPU is highly recommended)")

    if opts.model == 'resnet18':
        model = resnet18(pretrained=False)
        if opts.batchsize == -1:
            opts.batchsize = 1024 if opts.half else 512
    elif opts.model == 'resnet50':
        model = resnet50(pretrained=False)
        if opts.batchsize == -1:
            opts.batchsize = 512 if opts.half else 256
    elif opts.model == 'resnet101':
        model = torch.hub.load('pytorch/vision:v0.6.0', 'resnet101', pretrained=False)
        if opts.batchsize == -1:
            opts.batchsize = 360 if opts.half else 180
    elif opts.model == 'resnet152':
        model = torch.hub.load('pytorch/vision:v0.6.0', 'resnet152', pretrained=False)
        if opts.batchsize == -1:
            opts.batchsize = 256 if opts.half else 128
    elif opts.model == 'ran56':
         model = ResidualAttentionModel_56()
         if opts.batchsize == -1:
            opts.batchsize = 140
    elif opts.model == 'ran92':
         model = ResidualAttentionModel_92()
         if opts.batchsize == -1:
            opts.batchsize = 80      
    

    ch = model.fc.in_features
    model.fc = nn.Sequential(nn.Linear(ch, ch), nn.ReLU(), nn.Linear(ch, ch))

    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    # model = torch.nn.DataParallel(model)
    model.eval()
    if opts.half:
        model.half()
    for layer in model.modules():
        if isinstance(layer, nn.BatchNorm2d):
            layer.float()

    parameters = filter(lambda p: p.requires_grad, model.parameters())
    n_parameters = sum([p.data.nelement() for p in model.parameters()])
    print('  + Number of params: {}'.format(n_parameters))


    ### GPU Setup ###
    if use_gpu:
        model.cuda()
    ### DO NOT MODIFY THIS BLOCK ###
    if IS_ON_NSML:
        bind_nsml(model)
        if opts.pause:
            nsml.paused(scope=locals())
    ################################
    
    bind_nsml(model)
    
    if opts.mode == 'train':
        model.train()
        # Set dataloader
        train_ids, val_ids, unl_ids = split_ids(os.path.join(DATASET_PATH, 'train/train_label'), 0.2)
        print('found {} train, {} validation and {} unlabeled images'.format(len(train_ids), len(val_ids), len(unl_ids)))
        label_loader = torch.utils.data.DataLoader(
            SimpleImageLoader(DATASET_PATH, 'train', train_ids,
                              transform=transforms.Compose([
                                  transforms.Resize(opts.imResize),
                                  transforms.RandomResizedCrop(opts.imsize),
                                  # transforms.Resize((opts.imsize, opts.imsize)),
                                  transforms.RandomHorizontalFlip(),
                                  # transforms.ColorJitter(0.5, 0.5, 0.5, 0.5),
                                  transforms.ToTensor(),
                                  transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),])),
                                batch_size=opts.batchsize * 2, shuffle=True, num_workers=4, pin_memory=True, drop_last=True)
        print('train_loader done')

        unlabel_loader = torch.utils.data.DataLoader(
            SimpleImageLoader(DATASET_PATH, 'unlabel', unl_ids,
                              transform=transforms.Compose([
                                  transforms.Resize(opts.imResize),
                                  transforms.RandomResizedCrop(opts.imsize),
                                  transforms.ColorJitter(0.5, 0.5, 0.5, 0.5),
                                  transforms.ToTensor(),
                                  transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),])),
                                batch_size=opts.batchsize, shuffle=True, num_workers=0, pin_memory=True, drop_last=True)
        print('unlabel_loader done')

        validation_loader = torch.utils.data.DataLoader(
            SimpleImageLoader(DATASET_PATH, 'val', val_ids,
                               transform=transforms.Compose([
                                   transforms.Resize(opts.imResize),
                                   transforms.CenterCrop(opts.imsize),
                                   transforms.ToTensor(),
                                   transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),])),
                               batch_size=opts.batchsize, shuffle=False, num_workers=0, pin_memory=True, drop_last=False)
        print('validation_loader done')

        model = torch.nn.DataParallel(model) 
        model.to(device)
        bind_nsml(model)
        #Set optimizer
        if opts.optimizer == 'SGD':
            if opts.adaptive_lr:
                base_optimizer = optim.SGD(model.parameters(), lr=0.3*opts.batchsize/256)
            else:
                if opts.lr == -1:
                    base_optimizer = optim.SGD(model.parameters(), lr=0.001)
                else:
                    base_optimizer = optim.SGD(model.parameters(), lr=opts.lr)

        elif opts.optimizer == 'Adam':
            if opts.adaptive_lr:
                base_optimizer = optim.Adam(model.parameters(), lr=0.3*opts.batchsize/256)
            else:
                if opts.lr == -1:
                    base_optimizer = optim.Adam(model.parameters(), lr=0.001)
                else:
                    base_optimizer = optim.Adam(model.parameters(), lr=opts.lr)
        if opts.LARS:
            optimizer = torchlars.LARS(optimizer=base_optimizer, eps=1e-8, trust_coef=0.001)
        else: 
            optimizer = base_optimizer

        # INSTANTIATE LOSS CLASS
        unlabel_criterion = nt_cross_entropy

        # INSTANTIATE STEP LEARNING SCHEDULER CLASS
        if opts.scheduler == 'linear':
            scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer,  milestones=[50, 150], gamma=0.1)
        elif opts.scheduler == 'exp':
            scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=1e-6)

        model.train()
        print('==================================')
        print(opts)
        print('==================================')
        print('starting pretask')
        total_iter = 0
        for epoch in range(1, 201):
            for it, data in enumerate(unlabel_loader):
                total_iter += 1
                d = data.size()
                if opts.half:
                    x = data.view(d[0]*2, d[2], d[3], d[4]).half().to(device)
                else:
                    x = data.view(d[0]*2, d[2], d[3], d[4]).to(device)

                optimizer.zero_grad()
                p = model(x)
                if opts.half:
                    loss = unlabel_criterion(p.float())
                else: 
                    loss = unlabel_criterion(p)

                loss.backward()
                if opts.half:
                    model.float()
                optimizer.step()
                if opts.half:
                    model.half()
                    for layer in model.modules():
                        if isinstance(layer, nn.BatchNorm2d):
                            layer.float()
                print("epoch: ", epoch,  "loss: ", loss.item())
                nsml.report(summary=True, loss=loss.item(), step=total_iter)
            scheduler.step()
            print("epoch: ", epoch,  "loss: ", loss.item())
            if (epoch) % 2 == 0:
                if IS_ON_NSML:
                    nsml.save(opts.name + '_pre{}'.format(epoch))
                else:
                    torch.save(model.state_dict(), os.path.join('runs', opts.name + '_pre{}'.format(epoch))) 
Ejemplo n.º 3
0
        labels = np.asarray(label_list)
        labels = np.concatenate((labels, labels), axis=0)

        y_train = keras.utils.to_categorical(labels, num_classes=num_classes)
        x_train = x_train.astype('float32')
        x_train /= 255
        print(len(labels), 'train samples')
        """ Callback """
        monitor = 'get_categorical_accuracy_keras'
        reduce_lr = ReduceLROnPlateau(monitor=monitor, patience=3)
        """ Training loop """
        for epoch in range(nb_epoch):
            res = model.fit(x_train,
                            y_train,
                            batch_size=batch_size,
                            initial_epoch=epoch,
                            epochs=epoch + 1,
                            callbacks=[reduce_lr],
                            verbose=1,
                            shuffle=True)
            print(res.history)
            train_loss, train_acc = res.history['loss'][0], res.history[
                'get_categorical_accuracy_keras'][0]
            nsml.report(summary=True,
                        epoch=epoch,
                        epoch_total=nb_epoch,
                        loss=train_loss,
                        acc=train_acc)
            nsml.save(epoch)
Ejemplo n.º 4
0
def main():

    global char2index
    global index2char
    global SOS_token
    global EOS_token
    global PAD_token

    parser = argparse.ArgumentParser(description='Speech hackathon Baseline')
    parser.add_argument('--hidden_size',
                        type=int,
                        default=512,
                        help='hidden size of model (default: 256)')
    parser.add_argument('--layer_size',
                        type=int,
                        default=3,
                        help='number of layers of model (default: 3)')
    parser.add_argument(
        '--bidirectional',
        action='store_true',
        help='use bidirectional RNN for encoder (default: False)')
    parser.add_argument(
        '--use_attention',
        action='store_true',
        help='use attention between encoder-decoder (default: False)')
    parser.add_argument('--batch_size',
                        type=int,
                        default=32,
                        help='batch size in training (default: 32)')
    parser.add_argument(
        '--workers',
        type=int,
        default=4,
        help='number of workers in dataset loader (default: 4)')
    parser.add_argument('--max_epochs',
                        type=int,
                        default=10,
                        help='number of max epochs in training (default: 10)')
    parser.add_argument('--lr',
                        type=float,
                        default=1e-04,
                        help='learning rate (default: 0.0001)')
    parser.add_argument('--teacher_forcing',
                        type=float,
                        default=0.5,
                        help='teacher forcing ratio in decoder (default: 0.5)')
    parser.add_argument('--max_len',
                        type=int,
                        default=80,
                        help='maximum characters of sentence (default: 80)')
    parser.add_argument('--no_cuda',
                        action='store_true',
                        default=False,
                        help='disables CUDA training')
    parser.add_argument('--seed',
                        type=int,
                        default=1,
                        help='random seed (default: 1)')
    parser.add_argument('--save_name',
                        type=str,
                        default='model',
                        help='the name of model in nsml or local')
    parser.add_argument('--mode', type=str, default='train')
    parser.add_argument("--pause", type=int, default=0)
    parser.add_argument("--visdom", type=bool, default=False)
    parser.add_argument("--use_stft",
                        type=bool,
                        default=False,
                        help="use stft or log mel + specaugmentation")
    parser.add_argument("--mels", type=int, default=128)
    parser.add_argument("--use_rnn", type=bool, default=False)

    # Low Frame Rate (stacking and skipping frames)
    parser.add_argument('--LFR_m',
                        default=4,
                        type=int,
                        help='Low Frame Rate: number of frames to stack')
    parser.add_argument('--LFR_n',
                        default=3,
                        type=int,
                        help='Low Frame Rate: number of frames to skip')
    # EncoderTrans
    parser.add_argument('--n_layers_enc',
                        default=2,
                        type=int,
                        help='Number of encoder stacks')
    parser.add_argument('--n_head',
                        default=4,
                        type=int,
                        help='Number of Multi Head Attention (MHA)')
    parser.add_argument('--d_k', default=64, type=int, help='Dimension of key')
    parser.add_argument('--d_v',
                        default=64,
                        type=int,
                        help='Dimension of value')
    parser.add_argument('--d_model',
                        default=512,
                        type=int,
                        help='Dimension of model')
    parser.add_argument('--d_inner',
                        default=512,
                        type=int,
                        help='Dimension of inner')
    parser.add_argument('--dropout',
                        default=0.1,
                        type=float,
                        help='Dropout rate')
    parser.add_argument('--pe_maxlen',
                        default=5000,
                        type=int,
                        help='Positional Encoding max len')
    # Decoder Trans
    parser.add_argument('--d_word_vec',
                        default=512,
                        type=int,
                        help='Dim of decoder embedding')
    parser.add_argument('--n_layers_dec',
                        default=2,
                        type=int,
                        help='Number of decoder stacks')
    parser.add_argument('--tgt_emb_prj_weight_sharing',
                        default=1,
                        type=int,
                        help='share decoder embedding with decoder projection')
    # TransLoss
    parser.add_argument('--label_smoothing',
                        default=0.1,
                        type=float,
                        help='label smoothing')
    # Optimizer
    parser.add_argument('--k',
                        default=1.0,
                        type=float,
                        help='tunable scalar multiply to learning rate')
    parser.add_argument('--warmup_steps',
                        default=4000,
                        type=int,
                        help='warmup steps')
    args = parser.parse_args()

    char2index, index2char = label_loader.load_label('./data/hackathon.labels')
    SOS_token = char2index['<s>']
    EOS_token = char2index['</s>']
    PAD_token = char2index['_']

    # Setting seed
    random.seed(args.seed)
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed_all(args.seed)

    # Setting device
    args.cuda = not args.no_cuda and torch.cuda.is_available()
    device = torch.device('cuda' if args.cuda else 'cpu')

    # Feature extractor
    if args.use_stft:
        feature_size = N_FFT / 2 + 1
    else:
        feature_size = args.mels

    # Actual model
    if args.use_rnn:  # RNN structure

        # Define model
        enc = EncoderRNN(feature_size,
                         args.hidden_size,
                         input_dropout_p=args.dropout,
                         dropout_p=args.dropout,
                         n_layers=args.layer_size,
                         bidirectional=args.bidirectional,
                         rnn_cell='gru',
                         variable_lengths=False)
        dec = DecoderRNN(len(char2index),
                         args.max_len,
                         args.hidden_size * (2 if args.bidirectional else 1),
                         SOS_token,
                         EOS_token,
                         n_layers=args.layer_size,
                         rnn_cell='gru',
                         bidirectional=args.bidirectional,
                         input_dropout_p=args.dropout,
                         dropout_p=args.dropout,
                         use_attention=args.use_attention)
        model = Seq2seq(enc, dec)
        model.flatten_parameters()

        # Parameters initialization
        for param in model.parameters():
            param.data.uniform_(-0.08, 0.08)

        model = nn.DataParallel(model).to(device)

        optimizer = optim.Adam(model.module.parameters(), lr=args.lr)
        criterion = nn.CrossEntropyLoss(reduction='sum',
                                        ignore_index=PAD_token).to(device)

        bind_model(args, model, optimizer)

        if args.pause == 1:
            nsml.paused(scope=locals())

        if args.mode != "train":
            return

        data_list = os.path.join(DATASET_PATH, 'train_data', 'data_list.csv')
        wav_paths = list()
        script_paths = list()

        with open(data_list, 'r') as f:
            for line in f:
                # line: "aaa.wav,aaa.label"
                wav_path, script_path = line.strip().split(',')
                wav_paths.append(
                    os.path.join(DATASET_PATH, 'train_data', wav_path))
                script_paths.append(
                    os.path.join(DATASET_PATH, 'train_data', script_path))

        best_loss = 1e10
        best_cer = 1e10
        begin_epoch = 0

        # load all target scripts for reducing disk i/o
        target_path = os.path.join(DATASET_PATH, 'train_label')
        load_targets(target_path)

        train_batch_num, train_dataset_list, valid_dataset = split_dataset(
            args, wav_paths, script_paths, valid_ratio=0.05)

        logger.info('start')

        if args.visdom:
            train_visual = Visual(train_batch_num)
            eval_visual = Visual(1)

        train_begin = time.time()

        for epoch in range(begin_epoch, args.max_epochs):

            train_queue = queue.Queue(args.workers * 2)

            train_loader = MultiLoader(train_dataset_list, train_queue,
                                       args.batch_size, args.workers)
            train_loader.start()

            if args.visdom:
                train_loss, train_cer = trainRNN(model, train_batch_num,
                                                 train_queue, criterion,
                                                 optimizer, device,
                                                 train_begin, args.workers, 10,
                                                 args.teacher_forcing,
                                                 train_visual)
            else:
                train_loss, train_cer = trainRNN(model, train_batch_num,
                                                 train_queue, criterion,
                                                 optimizer, device,
                                                 train_begin, args.workers, 10,
                                                 args.teacher_forcing)

            logger.info('Epoch %d (Training) Loss %0.4f CER %0.4f' %
                        (epoch, train_loss, train_cer))

            train_loader.join()

            valid_queue = queue.Queue(args.workers * 2)
            valid_loader = BaseDataLoader(valid_dataset, valid_queue,
                                          args.batch_size, 0)
            valid_loader.start()

            if args.visdom:
                eval_loss, eval_cer = evaluateRNN(model, valid_loader,
                                                  valid_queue, criterion,
                                                  device, eval_visual)
            else:
                eval_loss, eval_cer = evaluateRNN(model, valid_loader,
                                                  valid_queue, criterion,
                                                  device)

            logger.info('Epoch %d (Evaluate) Loss %0.4f CER %0.4f' %
                        (epoch, eval_loss, eval_cer))

            valid_loader.join()

            nsml.report(False,
                        step=epoch,
                        train_epoch__loss=train_loss,
                        train_epoch__cer=train_cer,
                        eval__loss=eval_loss,
                        eval__cer=eval_cer)

            best_loss_model = (eval_loss < best_loss)
            best_cer_model = (eval_cer < best_cer)
            nsml.save(args.save_name)

            if best_loss_model:
                nsml.save('best_loss')
                best_loss = eval_loss
            if best_cer_model:
                nsml.save('best_cer')
                best_cer = eval_cer

    else:  # Transformer structure
        # Define model
        enc = EncoderTrans(feature_size,
                           args.n_layers_enc,
                           args.n_head,
                           args.d_k,
                           args.d_v,
                           args.d_model,
                           args.d_inner,
                           dropout=args.dropout,
                           pe_maxlen=args.pe_maxlen)
        dec = DecoderTrans(
            SOS_token,
            EOS_token,
            len(char2index),
            args.d_word_vec,
            args.n_layers_dec,
            args.n_head,
            args.d_k,
            args.d_v,
            args.d_model,
            args.d_inner,
            dropout=args.dropout,
            tgt_emb_prj_weight_sharing=args.tgt_emb_prj_weight_sharing,
            pe_maxlen=args.pe_maxlen)
        model = Transformer(enc, dec)

        # Parameter initialization
        for param in model.parameters():
            param.data.uniform_(-0.08, 0.08)
        model = nn.DataParallel(model).to(device)

        optimizer = TransformerOptimizer(
            torch.optim.Adam(model.parameters(), betas=(0.9, 0.98), eps=1e-09),
            args.k, args.d_model, args.warmup_steps)
        criterion = nn.CrossEntropyLoss(reduction='sum',
                                        ignore_index=PAD_token).to(device)
        bind_model(args, model, optimizer)

        if args.pause == 1:
            nsml.paused(scope=locals())

        if args.mode != "train":
            return

        data_list = os.path.join(DATASET_PATH, 'train_data', 'data_list.csv')
        wav_paths = list()
        script_paths = list()

        with open(data_list, 'r') as f:
            for line in f:
                # line: "aaa.wav,aaa.label"
                wav_path, script_path = line.strip().split(',')
                wav_paths.append(
                    os.path.join(DATASET_PATH, 'train_data', wav_path))
                script_paths.append(
                    os.path.join(DATASET_PATH, 'train_data', script_path))

        best_loss = 1e10
        best_cer = 1e10
        begin_epoch = 0

        # load all target scripts for reducing disk i/o
        target_path = os.path.join(DATASET_PATH, 'train_label')
        load_targets(target_path)

        train_batch_num, train_dataset_list, valid_dataset = split_dataset(
            args, wav_paths, script_paths, valid_ratio=0.05)

        logger.info('start')

        train_begin = time.time()

        for epoch in range(begin_epoch, args.max_epochs):

            train_queue = queue.Queue(args.workers * 2)

            train_loader = MultiLoader(train_dataset_list, train_queue,
                                       args.batch_size, args.workers)
            train_loader.start()

            if args.visdom:
                train_loss, train_cer = trainTrans(model, train_batch_num,
                                                   train_queue, criterion,
                                                   optimizer, device,
                                                   train_begin, args.workers,
                                                   10, args.teacher_forcing,
                                                   train_visual)
            else:
                train_loss, train_cer = trainTrans(
                    model,
                    train_batch_num,
                    train_queue,
                    criterion,
                    optimizer,
                    device,
                    train_begin,
                    args.workers,
                    10,
                    args.teacher_forcing,
                    label_smoothing=args.label_smoothing)

            logger.info('Epoch %d (Training) Loss %0.4f CER %0.4f' %
                        (epoch, train_loss, train_cer))

            train_loader.join()

            valid_queue = queue.Queue(args.workers * 2)
            valid_loader = BaseDataLoader(valid_dataset, valid_queue,
                                          args.batch_size, 0)
            valid_loader.start()

            if args.visdom:
                eval_loss, eval_cer = evaluateTrans(model, valid_loader,
                                                    valid_queue, criterion,
                                                    device, eval_visual)
            else:
                eval_loss, eval_cer = evaluateTrans(
                    model,
                    valid_loader,
                    valid_queue,
                    criterion,
                    device,
                    label_smoothing=args.label_smoothing)

            logger.info('Epoch %d (Evaluate) Loss %0.4f CER %0.4f' %
                        (epoch, eval_loss, eval_cer))

            valid_loader.join()

            nsml.report(False,
                        step=epoch,
                        train_epoch__loss=train_loss,
                        train_epoch__cer=train_cer,
                        eval__loss=eval_loss,
                        eval__cer=eval_cer)

            best_loss_model = (eval_loss < best_loss)
            best_cer_model = (eval_cer < best_cer)
            nsml.save(args.save_name)

            if best_loss_model:
                nsml.save('best_loss')
                best_loss = eval_loss
            if best_cer_model:
                nsml.save('best_cer')
                best_cer = eval_cer
Ejemplo n.º 5
0
def train_fine(opts, train_loader, model, criterion, optimizer, ema_optimizer,
               epoch, use_gpu, scheduler, is_mixsim):
    global global_step
    scaler = torch.cuda.amp.GradScaler()

    losses = AverageMeter()
    losses_x = AverageMeter()

    losses_curr = AverageMeter()
    losses_x_curr = AverageMeter()

    weight_scale = AverageMeter()
    acc_top1 = AverageMeter()
    acc_top5 = AverageMeter()

    model.train()

    # nCnt =0
    out = False
    local_step = 0
    while not out:
        labeled_train_iter = iter(train_loader)
        for batch_idx in range(len(train_loader)):
            try:
                data = labeled_train_iter.next()
                inputs_x, targets_x = data
            except:
                labeled_train_iter = iter(train_loader)
                data = labeled_train_iter.next()
                inputs_x, targets_x = data

            batch_size = inputs_x.size(0)
            # Transform label to one-hot
            classno = NUM_CLASSES
            targets_org = targets_x
            targets_x = torch.zeros(batch_size,
                                    classno).scatter_(1, targets_x.view(-1, 1),
                                                      1)

            if use_gpu:
                # Send input value to device 0, where first parameters of MixSim_Model is.
                if is_mixsim:
                    dev0 = 'cuda:{}'.format(opts.gpu_ids.split(',')[0])
                else:
                    dev0 = 'cuda'
                inputs_x, targets_x = inputs_x.to(dev0), targets_x.to(dev0)

            optimizer.zero_grad()

            with torch.cuda.amp.autocast():
                _, logits = model(inputs_x)

            loss = criterion(logits, targets_x, opts.temperature)

            losses.update(loss.item(), inputs_x.size(0))

            losses_curr.update(loss.item(), inputs_x.size(0))

            # compute gradient and do SGD step
            scaler.scale(loss).backward()
            scaler.step(optimizer)
            ema_optimizer.step()
            scaler.update()
            scheduler.step()

            with torch.no_grad():
                # compute guessed labels of unlabel samples
                _, pred_x1 = model(inputs_x)

            if IS_ON_NSML and global_step % opts.log_interval == 0:
                nsml.report(step=global_step, loss=losses_curr.avg)
                losses_curr.reset()

            acc_top1b = top_n_accuracy_score(targets_org.data.cpu().numpy(),
                                             pred_x1.data.cpu().numpy(),
                                             n=1) * 100
            acc_top5b = top_n_accuracy_score(targets_org.data.cpu().numpy(),
                                             pred_x1.data.cpu().numpy(),
                                             n=5) * 100
            acc_top1.update(torch.as_tensor(acc_top1b), inputs_x.size(0))
            acc_top5.update(torch.as_tensor(acc_top5b), inputs_x.size(0))

            local_step += 1
            global_step += 1

            if local_step >= opts.steps_per_epoch:
                out = True
                break

    return losses.avg, acc_top1.avg, acc_top5.avg
Ejemplo n.º 6
0
        for epoch in range(config.epochs):

            for i, (data,
                    labels) in enumerate(_batch_loader(dataset, config.batch)):

                _, c = sess.run([train, cost], feed_dict={x: data, y: labels})
                w_v, b_v, h, c, p, a = sess.run(
                    [w, b, hypothesis, cost, prediction, accuracy],
                    feed_dict={
                        x: data,
                        y: labels
                    })
            print('epoch:', epoch, "accuracy:", a)
            nsml.report(summary=True,
                        scope=locals(),
                        epoch=epoch,
                        accuracy=accuracy,
                        step=epoch)

            # save session
            if IS_ON_NSML:  # on nsml
                nsml.save(epoch)
            else:  # on local
                if epoch % 1000 == 0:
                    print(b_v)
                    save_path = saver.save(
                        sess, os.path.join('save2',
                                           str(epoch) + '.ckpt'))
                    print("Model saved in file: %s" % save_path)

    # local test debug mode
Ejemplo n.º 7
0
 def on_epoch_end(self, epoch, logs={}):
     nsml.report(summary=True, epoch=epoch, loss=logs.get('loss'), val_loss=logs.get('val_loss'))
     nsml.save(self.prefix +'_'+ str(self.seed)+'_' +str(epoch))
Ejemplo n.º 8
0
def main():

    global char2index
    global index2char
    global SOS_token
    global EOS_token
    global PAD_token

    parser = argparse.ArgumentParser(description='Speech hackathon Baseline')
    parser.add_argument('--hidden_size',
                        type=int,
                        default=512,
                        help='hidden size of model (default: 256)')
    parser.add_argument('--layer_size',
                        type=int,
                        default=3,
                        help='number of layers of model (default: 3)')
    parser.add_argument('--dropout',
                        type=float,
                        default=0.2,
                        help='dropout rate in training (default: 0.2)')
    parser.add_argument(
        '--bidirectional',
        action='store_true',
        help='use bidirectional RNN for encoder (default: False)')
    parser.add_argument(
        '--use_attention',
        action='store_true',
        help='use attention between encoder-decoder (default: False)')
    parser.add_argument('--batch_size',
                        type=int,
                        default=32,
                        help='batch size in training (default: 32)')
    parser.add_argument(
        '--workers',
        type=int,
        default=4,
        help='number of workers in dataset loader (default: 4)')
    parser.add_argument('--max_epochs',
                        type=int,
                        default=50,
                        help='number of max epochs in training (default: 10)')
    parser.add_argument('--lr',
                        type=float,
                        default=1e-04,
                        help='learning rate (default: 0.0001)')
    parser.add_argument('--teacher_forcing',
                        type=float,
                        default=0.5,
                        help='teacher forcing ratio in decoder (default: 0.5)')
    parser.add_argument('--max_len',
                        type=int,
                        default=80,
                        help='maximum characters of sentence (default: 80)')
    parser.add_argument('--no_cuda',
                        action='store_true',
                        default=False,
                        help='disables CUDA training')
    parser.add_argument('--seed',
                        type=int,
                        default=1,
                        help='random seed (default: 1)')
    parser.add_argument('--save_name',
                        type=str,
                        default='model',
                        help='the name of model in nsml or local')
    parser.add_argument('--mode', type=str, default='train')
    parser.add_argument("--pause", type=int, default=0)

    args = parser.parse_args()

    char2index, index2char = label_loader.load_label('./hackathon.labels')
    SOS_token = char2index['<s>']
    EOS_token = char2index['</s>']
    PAD_token = char2index['_']

    random.seed(args.seed)
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed_all(args.seed)

    args.cuda = not args.no_cuda and torch.cuda.is_available()
    device = torch.device('cuda' if args.cuda else 'cpu')

    # N_FFT: defined in loader.py
    #feature_size = N_FFT / 2 + 1
    #feature_size = 40
    feature_size = 45
    #Mel function이면 다르게 해야할듯?

    enc = EncoderRNN(feature_size,
                     args.hidden_size,
                     input_dropout_p=args.dropout,
                     dropout_p=args.dropout,
                     n_layers=args.layer_size,
                     bidirectional=args.bidirectional,
                     rnn_cell='gru',
                     variable_lengths=False)

    dec = DecoderRNN(len(char2index),
                     args.max_len,
                     args.hidden_size * (2 if args.bidirectional else 1),
                     SOS_token,
                     EOS_token,
                     n_layers=args.layer_size,
                     rnn_cell='gru',
                     bidirectional=args.bidirectional,
                     input_dropout_p=args.dropout,
                     dropout_p=args.dropout,
                     use_attention=args.use_attention)

    model = Seq2seq(enc, dec)
    model.flatten_parameters()

    for param in model.parameters():
        param.data.uniform_(-0.08, 0.08)

    model = nn.DataParallel(model).to(device)

    optimizer = optim.Adam(model.module.parameters(),
                           lr=args.lr,
                           weight_decay=1e-06)
    criterion = nn.CrossEntropyLoss(reduction='sum',
                                    ignore_index=PAD_token).to(device)

    bind_model(model, optimizer)

    if args.pause == 1:
        nsml.paused(scope=locals())

    if args.mode != "train":
        return
    '''
    nsml.load(checkpoint='best', session='team228/sr-hack-2019-dataset/414')
    nsml.save('best')
    '''

    data_list = os.path.join(DATASET_PATH, 'train_data', 'data_list.csv')
    wav_paths = list()
    script_paths = list()

    with open(data_list, 'r') as f:
        for line in f:
            # line: "aaa.wav,aaa.label"
            wav_path, script_path = line.strip().split(',')
            wav_paths.append(os.path.join(DATASET_PATH, 'train_data',
                                          wav_path))
            script_paths.append(
                os.path.join(DATASET_PATH, 'train_data', script_path))

    best_loss = 1e10
    begin_epoch = 0

    # load all target scripts for reducing disk i/o
    target_path = os.path.join(DATASET_PATH, 'train_label')
    load_targets(target_path)

    train_batch_num, train_dataset_list, valid_dataset = split_dataset(
        args, wav_paths, script_paths, valid_ratio=0.05)

    logger.info('start')

    train_begin = time.time()

    for epoch in range(begin_epoch, args.max_epochs):

        train_queue = queue.Queue(args.workers * 2)

        train_loader = MultiLoader(train_dataset_list, train_queue,
                                   args.batch_size, args.workers)
        train_loader.start()

        train_loss, train_cer = train(model, train_batch_num, train_queue,
                                      criterion, optimizer, device,
                                      train_begin, args.workers, 10,
                                      args.teacher_forcing)
        logger.info('Epoch %d (Training) Loss %0.4f CER %0.4f' %
                    (epoch, train_loss, train_cer))

        train_loader.join()

        valid_queue = queue.Queue(args.workers * 2)
        valid_loader = BaseDataLoader(valid_dataset, valid_queue,
                                      args.batch_size, 0)
        valid_loader.start()

        eval_loss, eval_cer = evaluate(model, valid_loader, valid_queue,
                                       criterion, device)
        logger.info('Epoch %d (Evaluate) Loss %0.4f CER %0.4f' %
                    (epoch, eval_loss, eval_cer))

        valid_loader.join()

        nsml.report(False,
                    step=epoch,
                    train_epoch__loss=train_loss,
                    train_epoch__cer=train_cer,
                    eval__loss=eval_loss,
                    eval__cer=eval_cer)

        best_model = (eval_loss < best_loss)
        nsml.save(args.save_name)

        if best_model:
            nsml.save('best')
            best_loss = eval_loss
Ejemplo n.º 9
0
        t0 = time.time()
        for epoch in range(nb_epoch):
            t1 = time.time()
            res = model.fit_generator(generator=train_generator,
                                      steps_per_epoch=STEP_SIZE_TRAIN,
                                      initial_epoch=epoch,
                                      validation_data=val_generator,
                                      epochs=epoch + 1,
                                      callbacks=[reduce_lr],
                                      verbose=1,
                                      shuffle=True,
                                      validation_steps=val_generator.n //
                                      batch_size)
            t2 = time.time()
            print(res.history)
            print('Training time for one epoch : %.1f' % ((t2 - t1)))
            train_loss, train_acc = res.history['loss'][0], res.history['acc'][
                0]
            val_loss, val_acc = res.history['val_loss'][0], res.history[
                'val_acc'][0]
            nsml.report(summary=True,
                        step=epoch,
                        epoch=epoch,
                        epoch_total=nb_epoch,
                        loss=train_loss,
                        acc=train_acc,
                        val_loss=val_loss,
                        val_acc=val_acc)
            nsml.save(epoch)
        print('Total training time : %.1f' % (time.time() - t0))
Ejemplo n.º 10
0
                loss = criterion(predictions, label_vars)

                optimizer.zero_grad()
                loss.backward()
                if grad_clip:
                    grad_norm = nn.utils.clip_grad_norm(model.parameters(), 1)
                    grad_norm = grad_norm
                else:
                    grad_norm = 0
                optimizer.step()
                avg_loss += loss.data[0]
                accumulate(model_run, model)
                print('batch', i, 'loss', loss.data[0], 'norm', grad_norm)
                nsml.report(summary=True, scope=locals(), epoch=epoch,
                            epoch_total=config.epochs, train__loss=float(loss.data[0]),
                            step=epoch * total_batch + i)
                #print('batch:', i, 'loss:', loss.data[0])
            print('epoch:', epoch, ' train_loss:', float(avg_loss/total_batch))

            if epoch in [4, 6]:
                optimizer.param_groups[0]['lr'] /= 10
                print('reduce learning rate')

            # nsml ps, 혹은 웹 상의 텐서보드에 나타나는 값을 리포트하는 함수입니다.
            #
            #nsml.report(summary=True, scope=locals(), epoch=epoch, epoch_total=config.epochs,
            #            train__loss=float(avg_loss/total_batch), step=epoch)
            # DONOTCHANGE (You can decide how often you want to save the model)
            nsml.save(epoch)
Ejemplo n.º 11
0
def train(*, env_id, num_env, hps, num_timesteps, seed):
    venv = VecFrameStack(
        make_atari_env(env_id,
                       num_env,
                       seed,
                       wrapper_kwargs=dict(),
                       start_index=num_env * MPI.COMM_WORLD.Get_rank(),
                       max_episode_steps=hps.pop('max_episode_steps')),
        hps.pop('frame_stack'))
    venv.score_multiple = 1
    venv.record_obs = False
    ob_space = venv.observation_space
    ac_space = venv.action_space
    gamma = hps.pop('gamma')
    policy = {'rnn': CnnGruPolicy, 'cnn': CnnPolicy}[hps.pop('policy')]

    agent = PpoAgent(
        scope='ppo',
        ob_space=ob_space,
        ac_space=ac_space,
        stochpol_fn=functools.partial(
            policy,
            scope='pol',
            ob_space=ob_space,
            ac_space=ac_space,
            update_ob_stats_independently_per_gpu=hps.pop(
                'update_ob_stats_independently_per_gpu'),
            proportion_of_exp_used_for_predictor_update=hps.pop(
                'proportion_of_exp_used_for_predictor_update'),
            exploration_type=hps.pop("exploration_type"),
            beta=hps.pop("beta"),
        ),
        gamma=gamma,
        gamma_ext=hps.pop('gamma_ext'),
        lam=hps.pop('lam'),
        nepochs=hps.pop('nepochs'),
        nminibatches=hps.pop('nminibatches'),
        lr=hps.pop('lr'),
        cliprange=0.1,
        nsteps=128,
        ent_coef=0.001,
        max_grad_norm=hps.pop('max_grad_norm'),
        use_news=hps.pop("use_news"),
        comm=MPI.COMM_WORLD if MPI.COMM_WORLD.Get_size() > 1 else None,
        update_ob_stats_every_step=hps.pop('update_ob_stats_every_step'),
        int_coeff=hps.pop('int_coeff'),
        ext_coeff=hps.pop('ext_coeff'),
        noise_type=hps.pop('noise_type'),
        noise_p=hps.pop('noise_p'),
        use_sched=hps.pop('use_sched'),
        num_env=num_env,
        exp_name=hps.pop('exp_name'),
    )
    agent.start_interaction([venv])
    if hps.pop('update_ob_stats_from_random_agent'):
        agent.collect_random_statistics(num_timesteps=128 * 50)
    assert len(hps) == 0, "Unused hyperparameters: %s" % list(hps.keys())

    counter = 0
    while True:
        info = agent.step()
        n_updates = 0

        if info['update']:
            logger.logkvs(info['update'])
            logger.dumpkvs()

            if NSML:
                n_updates = int(info['update']['n_updates'])
                nsml_dict = {
                    k: np.float64(v)
                    for k, v in info['update'].items()
                    if isinstance(v, Number)
                }
                nsml.report(step=n_updates, **nsml_dict)

            counter += 1
        #if n_updates >= 40*1000: # 40K updates
        #    break
        if agent.I.stats['tcount'] > num_timesteps:
            break

    agent.stop_interaction()
Ejemplo n.º 12
0
def main():

    global char2index
    global index2char
    global SOS_token
    global EOS_token
    global PAD_token

    parser = argparse.ArgumentParser(description='Speech hackathon Baseline')
    parser.add_argument('--hidden_size', type=int, default=512, help='hidden size of model (default: 512)')
    parser.add_argument('--layer_size', type=int, default=3, help='number of layers of model (default: 3)')
    parser.add_argument('--dropout', type=float, default=0.2, help='dropout rate in training (default: 0.2)')
    parser.add_argument('--bidirectional', action='store_true', help='use bidirectional RNN for encoder (default: False)')
    parser.add_argument('--use_attention', action='store_true', help='use attention between encoder-decoder (default: False)')
    parser.add_argument('--batch_size', type=int, default=32, help='batch size in training (default: 32)')
    parser.add_argument('--workers', type=int, default=4, help='number of workers in dataset loader (default: 4)')
    parser.add_argument('--max_epochs', type=int, default=10, help='number of max epochs in training (default: 10)')
    parser.add_argument('--lr', type=float, default=1e-04, help='learning rate (default: 0.0001)')
    parser.add_argument('--teacher_forcing', type=float, default=0.5, help='teacher forcing ratio in decoder (default: 0.5)')
    parser.add_argument('--max_len', type=int, default=80, help='maximum characters of sentence (default: 80)')
    parser.add_argument('--no_cuda', action='store_true', help='disables CUDA training')
    parser.add_argument('--seed', type=int, default=1, help='random seed (default: 1)')
    parser.add_argument('--save_name', type=str, default='model', help='the name of model in nsml or local')
    parser.add_argument('--mode', type=str, default='train')
    parser.add_argument('--pause', type=int, default=0)

    parser.add_argument('--log_dir', help='directory for logging, valid in local only')
    parser.add_argument('--patience', type=int, help='patience before early stopping (default to None)')
    parser.add_argument('--weight_decay', type=float, default=0, help='weight for L2 regularization')
    parser.add_argument('--save_from', type=int, default=0, help='starting epoch to save models')
    parser.add_argument('--load_ckpt', nargs=2, help='session and checkpoint to load')

    parser.add_argument('--transformer_encoder', action='store_true')
    parser.add_argument('--share_params', action='store_true')

    args = parser.parse_args()

    for name, value in args.__dict__.items():
        print('{}:\t{}'.format(name, value))
    print()

    if nsml.IS_ON_NSML:
        args.log_dir = None

    if args.log_dir is not None:
        if not osp.exists(args.log_dir):
            os.makedirs(args.log_dir)

        with open(osp.join(args.log_dir, 'args.txt'), 'w') as f:
            for name, value in args.__dict__.items():
                f.write('{}\t{}\n'.format(name, value))

    char2index, index2char = label_loader.load_label('./hackathon.labels')
    SOS_token = char2index['<s>']
    EOS_token = char2index['</s>']
    PAD_token = char2index['_']

    random.seed(args.seed)
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed_all(args.seed)

    args.cuda = not args.no_cuda and torch.cuda.is_available()
    device = torch.device('cuda' if args.cuda else 'cpu')

    # N_FFT: defined in loader.py
    feature_size = N_FFT / 2 + 1

    if args.transformer_encoder:
        enc = Encoder(len_max_seq=1248, d_word_vec=257, n_layers=6, n_head=8, d_k=64, d_v=64,
                      d_model=257, d_inner=2048, dropout=0.1, share_params=args.share_params)
    else:
        enc = EncoderRNN(
            feature_size, args.hidden_size, input_dropout_p=args.dropout, dropout_p=args.dropout,
            n_layers=args.layer_size, bidirectional=args.bidirectional, rnn_cell='gru',
            variable_lengths=False)

    dec = DecoderRNN(
        len(char2index), args.max_len, args.hidden_size * (2 if args.bidirectional else 1),
        SOS_token, EOS_token, n_layers=args.layer_size, rnn_cell='gru',
        bidirectional=args.bidirectional, input_dropout_p=args.dropout, dropout_p=args.dropout,
        use_attention=args.use_attention)

    if args.transformer_encoder:
        model = Seq2SeqTransformerEncoder(enc, dec)
    else:
        model = Seq2seq(enc, dec)
    model.flatten_parameters()

    for param in model.parameters():
        param.data.uniform_(-0.08, 0.08)

    model = nn.DataParallel(model).to(device)

    optimizer = optim.Adam(model.module.parameters(), lr=args.lr, weight_decay=args.weight_decay)
    criterion = nn.CrossEntropyLoss(reduction='sum', ignore_index=PAD_token).to(device)

    bind_model(model, optimizer)

    if args.load_ckpt is not None:
        nsml.load(session=args.load_ckpt[0], checkpoint=args.load_ckpt[1])

    if args.pause == 1:
        nsml.paused(scope=locals())

    if args.mode != "train":
        return

    data_list = osp.join(DATASET_PATH, 'train_data', 'data_list.csv')
    wav_paths = list()
    script_paths = list()

    with open(data_list, 'r') as f:
        for line in f:
            # line: "aaa.wav,aaa.label"

            wav_path, script_path = line.strip().split(',')
            wav_paths.append(osp.join(DATASET_PATH, 'train_data', wav_path))
            script_paths.append(osp.join(DATASET_PATH, 'train_data', script_path))

    cnt_converged = 0
    best_loss = 1e10
    begin_epoch = 0

    # load all target scripts for reducing disk i/o
    target_path = osp.join(DATASET_PATH, 'train_label')
    load_targets(target_path)

    train_batch_num, train_dataset_list, valid_dataset = split_dataset(
        args, wav_paths, script_paths, valid_ratio=0.05)

    logger.info('start')

    train_begin = time.time()

    if args.log_dir is not None:
        train_writer = SummaryWriter(logdir=osp.join(args.log_dir, 'train'))
        valid_writer = SummaryWriter(logdir=osp.join(args.log_dir, 'valid'))
    else:
        train_writer, valid_writer = None, None

    for epoch in range(begin_epoch, args.max_epochs):
        if args.load_ckpt is not None:
            valid_queue = queue.Queue(args.workers * 2)
            valid_loader = BaseDataLoader(valid_dataset, valid_queue, args.batch_size, 0)
            valid_loader.start()

            eval_loss, eval_cer = evaluate(model, valid_loader, valid_queue, criterion, device)
            logger.info('Eval right after model loading (just for checking)')
            logger.info('Epoch %d (Evaluate) Loss %0.4f CER %0.4f' % (epoch, eval_loss, eval_cer))

        train_queue = queue.Queue(args.workers * 2)
        train_loader = MultiLoader(train_dataset_list, train_queue, args.batch_size, args.workers)
        train_loader.start()

        train_loss, train_cer = train(model, train_batch_num, train_queue, criterion, optimizer,
                                      device, train_begin, args.workers, 100, args.teacher_forcing)
        logger.info('Epoch %d (Training) Loss %0.4f CER %0.4f' % (epoch, train_loss, train_cer))
        if args.log_dir is not None:
            train_writer.add_scalar('epoch/loss', train_loss, epoch)
            train_writer.add_scalar('epoch/CER', train_cer, epoch)

        train_loader.join()

        valid_queue = queue.Queue(args.workers * 2)
        valid_loader = BaseDataLoader(valid_dataset, valid_queue, args.batch_size, 0)
        valid_loader.start()

        eval_loss, eval_cer = evaluate(model, valid_loader, valid_queue, criterion, device)
        logger.info('Epoch %d (Evaluate) Loss %0.4f CER %0.4f' % (epoch, eval_loss, eval_cer))
        if args.log_dir is not None:
            valid_writer.add_scalar('epoch/loss', eval_loss, epoch)
            valid_writer.add_scalar('epoch/CER', eval_cer, epoch)

            with open(osp.join(args.log_dir, 'loss.txt'), 'a') as f:
                f.write('epoch: {}, train: {:.6f}, valid: {:.6f}\n'.format(epoch, train_loss, eval_loss))
            with open(osp.join(args.log_dir, 'CER.txt'), 'a') as f:
                f.write('epoch: {}, train: {:.6f}, valid: {:.6f}\n'.format(epoch, train_cer, eval_cer))

        valid_loader.join()

        nsml.report(False, step=epoch, train_epoch__loss=train_loss, train_epoch__cer=train_cer,
                    eval__loss=eval_loss, eval__cer=eval_cer)

        if epoch > args.save_from:
            nsml.save(args.save_name + '_e{}'.format(epoch))

        best_model = (eval_loss < best_loss)
        if best_model:
            nsml.save('best')
            best_loss = eval_loss

        if eval_loss > best_loss:
            cnt_converged += 1
            if args.patience is not None and cnt_converged > args.patience:
                break
        else:
            cnt_converged = 0
Ejemplo n.º 13
0
def main():

    global char2index
    global index2char
    global SOS_token
    global EOS_token
    global PAD_token

    parser = argparse.ArgumentParser(description='Speech hackathon Baseline')
    parser.add_argument('--hidden_size',
                        type=int,
                        default=512,
                        help='hidden size of model (default: 256)')
    parser.add_argument('--layer_size',
                        type=int,
                        default=3,
                        help='number of layers of model (default: 3)')
    parser.add_argument('--dropout',
                        type=float,
                        default=0.2,
                        help='dropout rate in training (default: 0.2)')
    parser.add_argument(
        '--bidirectional',
        action='store_true',
        help='use bidirectional RNN for encoder (default: False)')
    parser.add_argument(
        '--use_attention',
        action='store_true',
        help='use attention between encoder-decoder (default: False)')
    parser.add_argument('--batch_size',
                        type=int,
                        default=32,
                        help='batch size in training (default: 32)')
    parser.add_argument(
        '--workers',
        type=int,
        default=4,
        help='number of workers in dataset loader (default: 4)')
    parser.add_argument('--max_epochs',
                        type=int,
                        default=10,
                        help='number of max epochs in training (default: 10)')
    parser.add_argument('--lr',
                        type=float,
                        default=1e-04,
                        help='learning rate (default: 0.0001)')
    parser.add_argument('--teacher_forcing',
                        type=float,
                        default=0.5,
                        help='teacher forcing ratio in decoder (default: 0.5)')
    parser.add_argument('--max_len',
                        type=int,
                        default=80,
                        help='maximum characters of sentence (default: 80)')
    parser.add_argument('--no_cuda',
                        action='store_true',
                        default=False,
                        help='disables CUDA training')
    parser.add_argument('--seed',
                        type=int,
                        default=1,
                        help='random seed (default: 1)')
    parser.add_argument('--save_name',
                        type=str,
                        default='model',
                        help='the name of model in nsml or local')
    parser.add_argument('--mode', type=str, default='train')
    parser.add_argument("--pause", type=int, default=0)
    parser.add_argument('--rnn_cell', type=str, default='gru')
    parser.add_argument("--iteration", type=int, default=0)
    parser.add_argument('--feature', type=str, default='spec')
    parser.add_argument('--save_dir', type=str, default='')

    args = parser.parse_args()

    char2index, index2char = label_loader.load_label('./hackathon.labels')
    SOS_token = char2index['<s>']
    EOS_token = char2index['</s>']
    PAD_token = char2index['_']

    random.seed(args.seed)
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed_all(args.seed)

    args.cuda = not args.no_cuda and torch.cuda.is_available()
    device = torch.device('cuda' if args.cuda else 'cpu')

    logger.info('Using %s as feature' % args.feature)
    if args.save_dir:
        logger.info('Save directory: %s' % args.save_dir)
        os.makedirs(args.save_dir, exist_ok=True)

    # N_FFT: defined in loader.py
    if args.feature == 'mfcc':
        feature_size = N_MFCC * 3  # concat of mfcc, mfcc' mfcc''
    elif args.feature == 'melspec':
        feature_size = N_MELS
    elif args.feature == 'spec':
        feature_size = N_FFT / 2 + 1
    else:
        raise ValueError('Unsupported feature %s' % args.feature)

    enc = EncoderRNN(feature_size,
                     args.hidden_size,
                     input_dropout_p=args.dropout,
                     dropout_p=args.dropout,
                     n_layers=args.layer_size,
                     bidirectional=args.bidirectional,
                     rnn_cell=args.rnn_cell,
                     variable_lengths=False)

    dec = DecoderRNN(len(char2index),
                     args.max_len,
                     args.hidden_size * (2 if args.bidirectional else 1),
                     SOS_token,
                     EOS_token,
                     n_layers=args.layer_size,
                     rnn_cell=args.rnn_cell,
                     bidirectional=args.bidirectional,
                     input_dropout_p=args.dropout,
                     dropout_p=args.dropout,
                     use_attention=args.use_attention)

    model = Seq2seq(enc, dec)
    model.flatten_parameters()

    for param in model.parameters():
        param.data.uniform_(-0.08, 0.08)

    model = nn.DataParallel(model).to(device)

    optimizer = optim.Adam(model.module.parameters(), lr=args.lr)
    criterion = nn.CrossEntropyLoss(reduction='sum',
                                    ignore_index=PAD_token).to(device)

    bind_model(model, optimizer, args.feature)
    if args.pause != 1:
        nsml.load(checkpoint='10', session='team236/sr-hack-2019-dataset/122')
        nsml.save('init')
        logger.info('Saved!')

    if args.pause == 1:
        nsml.paused(scope=locals())

    if args.mode != "train":
        return

    data_list = os.path.join(DATASET_PATH, 'train_data', 'data_list.csv')
    wav_paths = list()
    script_paths = list()
    # load all target scripts for reducing disk i/o
    target_path = os.path.join(DATASET_PATH, 'train_label')
    target_dict = load_targets(target_path)

    with open(data_list, 'r') as f:
        for line in f:
            # line: "aaa.wav,aaa.label"
            wav_path, script_path = line.strip().split(',')
            wav_paths.append(os.path.join(DATASET_PATH, 'train_data',
                                          wav_path))
            script_paths.append(
                os.path.join(DATASET_PATH, 'train_data', script_path))

    best_loss = 1e10
    begin_epoch = 0

    train_dataset, valid_dataset = split_dataset(args,
                                                 wav_paths,
                                                 script_paths,
                                                 target_dict,
                                                 args.feature,
                                                 valid_ratio=0.05)

    train_begin = time.time()

    for epoch in range(begin_epoch, args.max_epochs):
        train_loader = torch.utils.data.DataLoader(train_dataset,
                                                   batch_size=args.batch_size,
                                                   shuffle=True,
                                                   num_workers=args.workers,
                                                   collate_fn=collate_fn)

        train_loss, train_cer = train(model, train_loader, criterion,
                                      optimizer, device, train_begin, 10,
                                      args.teacher_forcing)

        logger.info('Epoch %d (Training) Loss %0.4f CER %0.4f' %
                    (epoch, train_loss, train_cer))

        valid_loader = torch.utils.data.DataLoader(valid_dataset,
                                                   batch_size=4,
                                                   shuffle=False,
                                                   num_workers=args.workers,
                                                   collate_fn=collate_fn)

        eval_loss, eval_cer = evaluate(model, valid_loader, criterion, device)
        logger.info('Epoch %d (Evaluate) Loss %0.4f CER %0.4f' %
                    (epoch, eval_loss, eval_cer))

        nsml.report(False,
                    step=epoch,
                    train_epoch__loss=train_loss,
                    train_epoch__cer=train_cer,
                    eval__loss=eval_loss,
                    eval__cer=eval_cer)

        best_model = (eval_loss < best_loss)
        nsml.save(args.save_name)
        nsml.save(str(epoch))

        if args.save_dir:
            save_model(
                model, optimizer,
                os.path.join(args.save_dir,
                             './epoch-%d-cer-%d.pt' % (epoch, eval_cer)))

        if best_model:
            nsml.save('best')
            best_loss = eval_loss
Ejemplo n.º 14
0
def train(model,
          data_loader,
          criterion,
          optimizer,
          device,
          train_begin,
          print_batch=5,
          teacher_forcing_ratio=1):
    total_loss = 0.
    total_num = 0
    total_dist = 0
    total_length = 0
    total_sent_num = 0
    batch = 0

    model.train()

    total_batch_size = len(data_loader)
    logger.info('train() start')

    begin = epoch_begin = time.time()

    for batch_index, batch in enumerate(data_loader):
        feats, scripts, feat_lengths, script_lengths = batch
        optimizer.zero_grad()

        feats = feats.to(device)
        scripts = scripts.to(device)

        src_len = scripts.size(1)
        target = scripts[:, 1:]

        model.module.flatten_parameters()
        logit = model(feats,
                      feat_lengths,
                      scripts,
                      teacher_forcing_ratio=teacher_forcing_ratio)

        logit = torch.stack(logit, dim=1).to(device)

        y_hat = logit.max(-1)[1]

        loss = criterion(logit.contiguous().view(-1, logit.size(-1)),
                         target.contiguous().view(-1))
        total_loss += loss.item()
        total_num += sum(feat_lengths)

        display = random.randrange(0, 100) == 0
        dist, length = get_distance(target, y_hat, display=display)
        total_dist += dist
        total_length += length

        total_sent_num += target.size(0)

        loss.backward()
        optimizer.step()

        if batch_index % print_batch == 0:
            current = time.time()
            elapsed = current - begin
            epoch_elapsed = (current - epoch_begin) / 60.0
            train_elapsed = (current - train_begin) / 3600.0

            logger.info(
                'batch: {:4d}/{:4d}, loss: {:.4f}, cer: {:.2f}, elapsed: {:.2f}s {:.2f}m {:.2f}h'
                .format(batch_index, total_batch_size, total_loss / total_num,
                        total_dist / total_length, elapsed, epoch_elapsed,
                        train_elapsed))
            begin = time.time()

            nsml.report(False,
                        step=train.cumulative_batch_count,
                        train_step__loss=total_loss / total_num,
                        train_step__cer=total_dist / total_length)

        train.cumulative_batch_count += 1

    logger.info('train() completed')
    return total_loss / total_num, total_dist / total_length
Ejemplo n.º 15
0
    def train(self):
        with tf.control_dependencies(self.update_ops):
            step = tf.Variable(0, trainable=False)
            lr_ = config.lr
            lr = tf.placeholder(tf.float32)
            optimizer = tf.train.AdamOptimizer(lr, beta1=config.Beta1, beta2=config.Beta2)
            if (config.loss_seq in ('two', 'first', 'second')) & (config.hybrid == True):
                self.trainer_pre = optimizer.minimize(
                    self.loss_pre, var_list=tf.trainable_variables(), global_step=step)
            self.trainer = optimizer.minimize(
                self.loss, var_list=tf.trainable_variables(), global_step=step)

        NUM_THREADS = 3
        configuration = tf.ConfigProto(inter_op_parallelism_threads=NUM_THREADS, \
                                       intra_op_parallelism_threads=NUM_THREADS, \
                                       allow_soft_placement=True, \
                                       device_count={'CPU': 3}, \
                                       )

        with tf.Session(config=configuration) as sess:

            path = os.path.join(os.getcwd(), 'Results', 'Spectrograms')
            os.makedirs(path, exist_ok=True)
            path = os.path.join(os.getcwd(), 'Results', 'Test')
            os.makedirs(path, exist_ok=True)
            path = os.path.join(os.getcwd(), 'Results', 'Simulation')
            os.makedirs(path, exist_ok=True)
            path = os.path.join(os.getcwd(), 'Results', 'CheckPoint')
            os.makedirs(path, exist_ok=True)
            path = os.path.join(os.getcwd(), 'Results', 'Logs_Training')
            os.makedirs(path, exist_ok=True)

            tf.global_variables_initializer().run()

            coord = tf.train.Coordinator()
            threads = tf.train.start_queue_runners(sess, coord=coord)

            saver = tf.train.Saver(max_to_keep=10)
            writer = tf.summary.FileWriter('./Results/Logs_Training', sess.graph)

            """
            Restore
                Restore checkpoint only when the user specify the 'config.restore' to be True
            """
            if config.restore:
                ckpt = tf.train.get_checkpoint_state(checkpoint_dir="./Results/CheckPoint")
                try:
                    if ckpt and ckpt.model_checkpoint_path:
                        print("check point path : ", ckpt.model_checkpoint_path)
                        saver.restore(sess, ckpt.model_checkpoint_path)
                        print('Restored!')
                except AttributeError:
                    print("No checkpoint")
            else:
                print("No model restoration")

            training_loss = 0
            training_alpha_sum = 0
            training_alpha_list = []
            f = open(os.path.join(os.getcwd(), 'Results', 'Logs_Training', 'training_loss.txt'), 'a+')
            f.write("Variable Size : {}\n".format(self.variable_size))
            f.close()

            for iter in range(config.iter + 1):
                if coord.should_stop():
                    print("break data pipeline")
                    break

                if (config.loss_seq in ('two', 'first', 'second')) & (config.hybrid == True):
                    if iter <= config.loss_thereshold:
                        _, loss_, alpha_ = sess.run([self.trainer_pre, self.loss_pre, self.graph.alpha], feed_dict={lr: lr_})
                    else:
                        _, loss_, alpha_ = sess.run([self.trainer, self.loss, self.graph.alpha], feed_dict={lr: lr_})
                else:
                    _, loss_, alpha_ = sess.run([self.trainer, self.loss, self.graph.alpha], feed_dict={lr: lr_})

                training_loss += loss_
                training_alpha_sum += np.mean(alpha_)
                training_alpha_list.append(np.squeeze(alpha_))

                if iter % 10 == 0:
                    print("iter : {} (loss : {:.3f}, alpha : {:.3f}/ {:.3f})".format(iter, loss_, np.mean(alpha_), np.std(alpha_)))
                    summary_ = sess.run(self.summary)
                    writer.add_summary(summary_, iter)

                # print spectrograms
                if iter % 5000 == 0:
                    spec_mix_, spec_tgt_, spec_est_ = sess.run(
                        [self.graph.input_spec, self.graph.source_spec, self.graph.masked_spec])

                    mag_mix_ = np.sqrt(np.square(spec_mix_[:, :, :, 0]) + np.square(spec_mix_[:, :, :, 1]))
                    mag_tgt_ = np.sqrt(np.square(spec_tgt_[:, :, :, 0]) + np.square(spec_tgt_[:, :, :, 1]))
                    mag_est_ = np.sqrt(np.square(spec_est_[:, :, :, 0]) + np.square(spec_est_[:, :, :, 1]))

                    mag_image1 = np.log10(np.concatenate((mag_mix_[0], mag_mix_[1], mag_mix_[2]), axis=1) + 1e-02)
                    mag_image2 = np.log10(np.concatenate((mag_tgt_[0], mag_tgt_[1], mag_tgt_[2]), axis=1) + 1e-02)
                    mag_image3 = np.log10(np.concatenate((mag_est_[0], mag_est_[1], mag_est_[2]), axis=1) + 1e-02)

                    mag_image4 = np.log10(
                        np.concatenate((mag_mix_[-3], mag_mix_[-2], mag_mix_[-1]), axis=1) + 1e-02)
                    mag_image5 = np.log10(
                        np.concatenate((mag_tgt_[-3], mag_tgt_[-2], mag_tgt_[-1]), axis=1) + 1e-02)
                    mag_image6 = np.log10(
                        np.concatenate((mag_est_[-3], mag_est_[-2], mag_est_[-1]), axis=1) + 1e-02)

                    fig = plt.figure(figsize=(15, 15))
                    fig.suptitle('Spectrograms', fontsize=20, family='serif')

                    font = {'family': 'serif',
                            'color': 'darkred',
                            'weight': 'normal',
                            'size': 15,
                            }

                    ax1 = fig.add_subplot(1, 3, 1)
                    ax1.set_title('Mixture', fontdict=font)
                    ax1.imshow(mag_image1, interpolation='nearest', aspect='auto', cmap='jet')
                    ax1.xaxis.set_tick_params(labelsize=15)
                    ax1.yaxis.set_tick_params(labelsize=15)
                    plt.gca().invert_yaxis()

                    ax2 = fig.add_subplot(1, 3, 2)
                    ax2.set_title('True Vocal', fontdict=font)
                    ax2.imshow(mag_image2, interpolation='nearest', aspect='auto', cmap='jet')
                    ax2.xaxis.set_tick_params(labelsize=15)
                    ax2.yaxis.set_tick_params(labelsize=15)
                    plt.gca().invert_yaxis()

                    ax3 = fig.add_subplot(1, 3, 3)
                    ax3.set_title('Estimated Vocal', fontdict=font)
                    ax3.imshow(mag_image3, interpolation='nearest', aspect='auto', cmap='jet')
                    ax3.xaxis.set_tick_params(labelsize=15)
                    ax3.yaxis.set_tick_params(labelsize=15)
                    plt.gca().invert_yaxis()

                    plt.tight_layout(pad=0.4, w_pad=1, h_pad=1.0)
                    plt.subplots_adjust(top=0.85)
                    fig.savefig("./Results/Spectrograms/Spectrogram{}.png".format(iter))
                    plt.close(fig)

                    fig = plt.figure(figsize=(15, 15))
                    fig.suptitle('NoiseSpectrograms', fontsize=20, family='serif')

                    font = {'family': 'serif',
                            'color': 'darkred',
                            'weight': 'normal',
                            'size': 15,
                            }

                    ax1 = fig.add_subplot(1, 3, 1)
                    ax1.set_title('Noise', fontdict=font)
                    ax1.imshow(mag_image4, interpolation='nearest', aspect='auto', cmap='jet')
                    ax1.xaxis.set_tick_params(labelsize=15)
                    ax1.yaxis.set_tick_params(labelsize=15)
                    plt.gca().invert_yaxis()

                    ax2 = fig.add_subplot(1, 3, 2)
                    ax2.set_title('True Zeros', fontdict=font)
                    ax2.imshow(mag_image5, interpolation='nearest', aspect='auto', cmap='jet')
                    ax2.xaxis.set_tick_params(labelsize=15)
                    ax2.yaxis.set_tick_params(labelsize=15)
                    plt.gca().invert_yaxis()

                    ax3 = fig.add_subplot(1, 3, 3)
                    ax3.set_title('Estimated Zeros', fontdict=font)
                    ax3.imshow(mag_image6, interpolation='nearest', aspect='auto', cmap='jet')
                    ax3.xaxis.set_tick_params(labelsize=15)
                    ax3.yaxis.set_tick_params(labelsize=15)
                    plt.gca().invert_yaxis()

                    plt.tight_layout(pad=0.4, w_pad=1, h_pad=1.0)
                    plt.subplots_adjust(top=0.85)
                    fig.savefig(
                        "./Results/Spectrograms/Noise_input_Spectrogram{}.png".format(iter))
                    plt.close(fig)

                if (iter) % 10000 == 0:
                    training_alpha_std = np.std(np.array(training_alpha_list))
                    f = open(os.path.join(os.getcwd(), 'Results', 'Logs_Training', 'training_loss.txt'), 'a+')
                    f.write("average {} loss : {:.4f} (iter {} ~ {})\n".format(config.loss, training_loss / 10000,
                                                                               iter - 9999, iter + 1))
                    f.close()
                    f = open(os.path.join(os.getcwd(), 'Results', 'Logs_Training', 'training_alpha.txt'), 'a+')
                    f.write("average alpha : {:.4f} / std : {:.4f} (iter {} ~ {})\n".format(training_alpha_sum / 10000,
                                                                                            training_alpha_std ,iter - 9999, iter + 1))
                    f.close()

                    nsml.report(summary=True,
                                iter=iter,
                                train__loss=float(training_loss / 10000),
                                alpha=float(training_alpha_sum / 10000),
                                alpha_std=float(training_alpha_std),
                                lr=float(lr_)*10000
                                )
                    training_loss = 0
                    training_alpha_sum = 0
                    training_alpha_list = []

                if (iter+1) % config.lr_decay_iter == 0:
                    if config.lr_decay:
                        lr_ = lr_/config.lr_decay_rate
                        print("learning rate is decayed : {}".format(lr_))

                # test songs
                if (iter + 1) % config.model_save_iter == 0:
                    moving_rate = config.moving_rate
                    hop = int(config.audio_size / moving_rate)  # This has to be an interger without any rounding
                    mask = (signal.get_window('hanning', config.audio_size)) / (moving_rate / 2)
                    # mask = (1 / moving_rate)

                    for folder, root in [('Test', config.test_sample_root), ('Simulation', config.simul_sample_root)]:
                        test_sample_path = glob(os.path.join(root, "noisy/*.wav"))
                        print("{} audios for test".format(len(test_sample_path)))

                        SNR_list = []
                        infer_time = []

                        wb = openpyxl.Workbook()
                        ws = wb.active

                        for k, song_path in enumerate(sorted(test_sample_path)):
                            print("\n %d th song" % k)

                            # Load song
                            name = os.path.basename(song_path)
                            print("Loading {}...".format(name))
                            noisy, _ = librosa.core.load(song_path, sr=16000)
                            clean, _ = librosa.core.load(os.path.join(root, "spk", name[:8]+".wav"),
                                                         sr=16000)
                            noise = noisy - clean

                            print("song shape : ", noisy.shape)

                            song_length = noisy.shape[0]
                            num_frame = int(song_length / config.audio_size) + 1

                            # pad zeros to song to make the length of it be the multiple of config.audio_size
                            noisy_pad = np.lib.pad(noisy, (
                                config.audio_size, num_frame * config.audio_size - song_length + config.audio_size),
                                                   'constant')
                            clean_pad = np.lib.pad(clean, (
                                config.audio_size, num_frame * config.audio_size - song_length + config.audio_size),
                                                   'constant')
                            padded_song_shape = noisy_pad.shape[0]

                            # Slice mixture and source
                            noisy_slice_list = []
                            clean_slice_list = []

                            hop = int(
                                config.audio_size / config.moving_rate)  # This has to be an interger without any rounding
                            num_to_move = int((padded_song_shape - config.audio_size) / hop) + 1
                            for i in range(num_to_move):
                                start_point = int(i * hop)
                                end_point = int(start_point + config.audio_size)
                                noisy_slice = noisy_pad[start_point:end_point]  # ex: (1,16384)
                                noisy_slice_list.append(noisy_slice)
                                clean_slice = clean_pad[start_point:end_point]
                                clean_slice_list.append(clean_slice)

                            num_slice = len(noisy_slice_list)
                            noisy_stacked = np.array(noisy_slice_list)
                            clean_stacked = np.array(clean_slice_list)

                            # Separation
                            segments = []
                            spec_mix_list = []
                            spec_tgt_list = []
                            spec_est_list = []
                            spec_est1_list = []
                            spec_est2_list = []
                            alpha_list = []

                            for n, i in enumerate(range(num_slice)):
                                start_time = time.time()
                                print("processing {}: {}/{}".format(name, n, num_slice))
                                estimated_sample, spec_mix_, spec_tgt_, spec_est_, spec_est1_,  spec_est2_, alpha_ \
                                    = sess.run([self.graph_test.estimated, self.graph_test.input_spec,
                                                self.graph_test.source_spec, self.graph_test.masked_spec,
                                                self.graph_test.masked_spec1,  self.graph_test.masked_spec2,
                                                self.graph_test.alpha],
                                               feed_dict={self.mix_test: noisy_stacked[i:i + 1],
                                                          self.source_test: clean_stacked[i:i + 1]})
                                infer_time.append(time.time() - start_time)
                                masked_sample = np.expand_dims(mask, axis=0) * estimated_sample
                                segments.append(masked_sample)
                                alpha_list.append(np.mean(alpha_))

                                if (n+1) % 4 == 0:
                                    spec_mix_list.append(spec_mix_)
                                    spec_tgt_list.append(spec_tgt_)
                                    spec_est_list.append(spec_est_)
                                    spec_est1_list.append(spec_est1_)
                                    spec_est2_list.append(spec_est2_)

                            spec_mix_ = np.concatenate(spec_mix_list, axis=0)
                            spec_tgt_ = np.concatenate(spec_tgt_list, axis=0)
                            spec_est_ = np.concatenate(spec_est_list, axis=0)
                            spec_est1_ = np.concatenate(spec_est1_list, axis=0)
                            spec_est2_ = np.concatenate(spec_est2_list, axis=0)

                            # spectrograms
                            mag_mix_ = np.sqrt(np.square(spec_mix_[:, :, :, 0]) + np.square(spec_mix_[:, :, :, 1]))
                            mag_tgt_ = np.sqrt(np.square(spec_tgt_[:, :, :, 0]) + np.square(spec_tgt_[:, :, :, 1]))
                            mag_est_ = np.sqrt(np.square(spec_est_[:, :, :, 0]) + np.square(spec_est_[:, :, :, 1]))
                            mag_est1_ = np.sqrt(np.square(spec_est1_[:, :, :, 0]) + np.square(spec_est1_[:, :, :, 1]))
                            mag_est2_ = np.sqrt(np.square(spec_est2_[:, :, :, 0]) + np.square(spec_est2_[:, :, :, 1]))

                            mag_image1 = np.log10(np.concatenate((mag_mix_[0], mag_mix_[1], mag_mix_[2]), axis=1) + 1e-02)
                            mag_image2 = np.log10(np.concatenate((mag_tgt_[0], mag_tgt_[1], mag_tgt_[2]), axis=1) + 1e-02)
                            mag_image3 = np.log10(np.concatenate((mag_est_[0], mag_est_[1], mag_est_[2]), axis=1) + 1e-02)
                            mag_image4 = np.log10(np.concatenate((mag_est1_[0], mag_est1_[1], mag_est1_[2]), axis=1) + 1e-02)
                            mag_image5 = np.log10(np.concatenate((mag_est2_[0], mag_est1_[1], mag_est2_[2]), axis=1) + 1e-02)

                            fig = plt.figure(figsize=(25, 15))
                            fig.suptitle('Spectrograms', fontsize=20, family='serif')

                            font = {'family': 'serif',
                                    'color': 'darkred',
                                    'weight': 'normal',
                                    'size': 15,
                                    }

                            ax1 = fig.add_subplot(1, 5, 1)
                            ax1.set_title('Mixture', fontdict=font)
                            ax1.imshow(mag_image1, interpolation='nearest', aspect='auto', cmap='jet')
                            ax1.xaxis.set_tick_params(labelsize=15)
                            ax1.yaxis.set_tick_params(labelsize=15)
                            plt.gca().invert_yaxis()

                            ax2 = fig.add_subplot(1, 5, 2)
                            ax2.set_title('True Vocal', fontdict=font)
                            ax2.imshow(mag_image2, interpolation='nearest', aspect='auto', cmap='jet')
                            ax2.xaxis.set_tick_params(labelsize=15)
                            ax2.yaxis.set_tick_params(labelsize=15)
                            plt.gca().invert_yaxis()

                            ax3 = fig.add_subplot(1, 5, 3)
                            ax3.set_title('Estimated Vocal {:.3f}'.format(np.mean(alpha_)), fontdict=font)
                            ax3.imshow(mag_image3, interpolation='nearest', aspect='auto', cmap='jet')
                            ax3.xaxis.set_tick_params(labelsize=15)
                            ax3.yaxis.set_tick_params(labelsize=15)
                            plt.gca().invert_yaxis()

                            ax3 = fig.add_subplot(1, 5, 4)
                            ax3.set_title('Estimated-1 Vocal', fontdict=font)
                            ax3.imshow(mag_image4, interpolation='nearest', aspect='auto', cmap='jet')
                            ax3.xaxis.set_tick_params(labelsize=15)
                            ax3.yaxis.set_tick_params(labelsize=15)
                            plt.gca().invert_yaxis()

                            ax3 = fig.add_subplot(1, 5, 5)
                            ax3.set_title('Estimated-2 Vocal', fontdict=font)
                            ax3.imshow(mag_image5, interpolation='nearest', aspect='auto', cmap='jet')
                            ax3.xaxis.set_tick_params(labelsize=15)
                            ax3.yaxis.set_tick_params(labelsize=15)
                            plt.gca().invert_yaxis()

                            plt.tight_layout(pad=0.4, w_pad=1, h_pad=1.0)
                            plt.subplots_adjust(top=0.85)
                            fig.savefig("./Results/{}/{}_{}.png".format(folder, iter, os.path.splitext(name)[0]))
                            plt.close(fig)

                            # Post-processing(triangle mask)
                            # num_to_pad = int((config.audio_size/2)*(num_slice-1))
                            # temp = np.zeros(shape=(config.audio_size + num_to_pad))
                            temp = np.zeros(shape=(padded_song_shape))
                            for i in range(len(segments)):
                                start_point = int(i * (config.audio_size / config.moving_rate))
                                end_point = int(start_point + config.audio_size)
                                temp[start_point:end_point] = temp[start_point:end_point] + segments[i]

                            # Match the original song length
                            estimation = np.squeeze(temp[config.audio_size:config.audio_size + song_length])
                            # print("estimation shape: {}".format(estimation.shape))

                            # save separated source as audio
                            OUTPUT_FILENAME = "./Results/{}/{}_{}_{}_{}.wav".format(folder, iter, os.path.splitext(name)[0],
                                                                                  config.network, config.network2)
                            librosa.output.write_wav(OUTPUT_FILENAME, estimation, 16000)

                            s_target = clean * np.sum(estimation * clean) / np.sum(clean ** 2)
                            error = estimation - s_target
                            snr = 10 * np.log10(sum(s_target ** 2) / sum(error ** 2))
                            SNR_list.append(snr)

                            noise_type = os.path.splitext(name)[0].split("_")[2]
                            noise_snr = os.path.splitext(name)[0].split("_")[3]

                            ws.cell(row= 3, column= k % 3 + 2).value = noise_snr
                            ws.cell(row= k//3 + 4, column= k % 3 + 2).value = snr
                            ws.cell(row= k//3 + 4, column= 1).value = noise_type
                            ws.cell(row= 3, column= k % 3 + 6).value = noise_snr
                            ws.cell(row= k//3 + 4, column= k % 3 + 6).value = np.mean(np.array(alpha_list))

                        snr_mean = np.mean(np.array(SNR_list))
                        snr_std = np.std(np.array(SNR_list))
                        print("snr mean : {}, std : {}".format(snr_mean, snr_std))
                        ws.cell(row=1, column=2).value = "mean"
                        ws.cell(row=1, column=3).value = "std"
                        ws.cell(row=1, column=4).value = "infer_time per frame"
                        ws.cell(row=2, column=2).value = snr_mean
                        ws.cell(row=2, column=3).value = snr_std
                        ws.cell(row=2, column=4).value = np.mean(np.array(infer_time))
                        ws.cell(row=2, column=6).value = "alpha"

                        wb.save('./Results/{}/snr_{}_{}_{}.xlsx'.format(folder, iter, config.network, config.network2))
                        wb.close()

                    print("Save model for epoch {}".format(iter))
                    saver.save(sess, "./Results/CheckPoint/model.ckpt", global_step=(iter+1))
                    print("-------------------Model saved-------------------")

            coord.request_stop()
            coord.join(threads)
Ejemplo n.º 16
0
def train(model,
          total_batch_size,
          queue,
          criterion,
          optimizer,
          device,
          train_begin,
          train_loader_count,
          print_batch=5,
          teacher_forcing_ratio=1):
    total_loss = 0.
    total_num = 0
    total_dist = 0
    total_length = 0
    total_sent_num = 0
    batch = 0

    model.train()

    logger.info('train() start')

    begin = epoch_begin = time.time()

    while True:
        if queue.empty():
            logger.debug('queue is empty')

        feats, scripts, feat_lengths, script_lengths = queue.get()

        if feats.shape[0] == 0:
            # empty feats means closing one loader
            train_loader_count -= 1

            logger.debug('left train_loader: %d' % (train_loader_count))

            if train_loader_count == 0:
                break
            else:
                continue

        optimizer.zero_grad()

        feats = feats.to(device)
        scripts = scripts.to(device)

        src_len = scripts.size(1)
        target = scripts[:, 1:]

        model.module.flatten_parameters()
        logit = model(feats,
                      feat_lengths,
                      scripts,
                      teacher_forcing_ratio=teacher_forcing_ratio)

        logit = torch.stack(logit, dim=1).to(device)

        y_hat = logit.max(-1)[1]

        loss = criterion(logit.contiguous().view(-1, logit.size(-1)),
                         target.contiguous().view(-1))
        total_loss += loss.item()
        total_num += sum(feat_lengths)

        display = random.randrange(0, 100) == 0
        dist, length = get_distance(target, y_hat, display=display)
        total_dist += dist
        total_length += length

        total_sent_num += target.size(0)

        loss.backward()
        optimizer.step()

        if batch % print_batch == 0:
            current = time.time()
            elapsed = current - begin
            epoch_elapsed = (current - epoch_begin) / 60.0
            train_elapsed = (current - train_begin) / 3600.0

            logger.info(
                'batch: {:4d}/{:4d}, loss: {:.4f}, cer: {:.2f}, elapsed: {:.2f}s {:.2f}m {:.2f}h'
                .format(
                    batch,
                    #len(dataloader),
                    total_batch_size,
                    total_loss / total_num,
                    total_dist / total_length,
                    elapsed,
                    epoch_elapsed,
                    train_elapsed))
            begin = time.time()

            nsml.report(False,
                        step=train.cumulative_batch_count,
                        train_step__loss=total_loss / total_num,
                        train_step__cer=total_dist / total_length)
        batch += 1
        train.cumulative_batch_count += 1

    logger.info('train() completed')
    return total_loss / total_num, total_dist / total_length
Ejemplo n.º 17
0
 def on_epoch_end(self, epoch, logs={}):
     nsml.report(summary=True, scope=locals(), epoch=epoch, epoch_total=config.epochs, step=epoch)
     nsml.save(epoch)
Ejemplo n.º 18
0
                          ', Accuracy:', round(accuracy, 2), ', time: %.2f' % (time.time() - t0))
                    t0 = time.time()
                if (i+1) >= total_eval_batch:
                    break

            if total_eval_batch > 0:
                eval_accuracy = float(avg_accuracy / total_eval_batch)
                eval_loss = float(avg_loss/total_eval_batch)
                print('\t  eval_loss: %.3f' % eval_loss, 'accuracy: %.2f' % eval_accuracy,
                        ', time: %.1fs' % (time.time() - t1))
            else:
                eval_accuracy = 0
                eval_loss = 0

            # nsml ps, 혹은 웹 상의 텐서보드에 나타나는 값을 리포트하는 함수입니다.
            nsml.report(summary=True, scope=locals(), epoch=epoch, epoch_total=config.epochs,
                        train__loss=train_loss, train_accuracy=train_accuracy,
                        eval__loss=eval_loss, eval_accuracy=eval_accuracy,
                        step=epoch)
            # DONOTCHANGE (You can decide how often you want to save the model)
            nsml.save(epoch)

    # 로컬 테스트 모드일때 사용합니다
    # 결과가 아래와 같이 나온다면, nsml submit을 통해서 제출할 수 있습니다.
    # [(0.0, 9.045), (0.0, 5.91), ... ]
    elif config.mode == 'test_local':
        with open(os.path.join(DATASET_PATH, 'train/train_data'), 'rt', encoding='utf-8') as f:
            reviews = f.readlines()[:config.max_dataset]
        res = nsml.infer(reviews)
        print(res)
Ejemplo n.º 19
0
def train(model, train_loader, args):
    # Set mode to 'train' (for batch norm.)
    model.train()

    # Enable CUDA if specified
    use_cuda = args.cuda and torch.cuda.is_available()
    device = torch.device("cuda" if use_cuda else "cpu")
    model = model.to(device)

    # Initialize visdom object
    if not args.local:
        viz = Visdom(visdom=visdom)

    # Specify loss function & optimizer
    if args.loss == 'L1':
        criterion = nn.L1Loss()
    elif args.loss == 'MSE':
        criterion = nn.MSELoss()
    else:
        raise NotImplementedError("Only L1Loss(cmd arg:'L1') and MSELoss(cmd arg:'MSE') are supported.")
    parameters = [p for p in model.parameters() if p.requires_grad]
    optimizer = torch.optim.Adam(params = parameters, lr = args.learning_rate)

    # Compute the number of iterations in an epoch
    num_iters = len(train_loader)
    if args.max_dataset_size > 0:
        num_iters = args.max_dataset_size // args.batch_size

    min_loss = 10000000

    # For NSML saving function
    bind_model(model=model, optimizer=optimizer)

    # Add perceptual loss term (use upto 'relu2-2' layer of VGG-16)
    if args.gamma > 0:
        vgg16_model = models.vgg16_bn(pretrained=True).to(device)
        layers_list = list(vgg16_model.features.children())[:14]
        '''
        for idx in [2,5,9,12]:
            layers_list[idx] = nn.ReLU(inplace=False)
        for layer in layers_list:
            print(layer)
        '''
        perceptual_model = torch.nn.Sequential(*layers_list).to(device)
        for param in perceptual_model.parameters():
            param.requires_grad = False
        
        perceptual_criterion = nn.MSELoss()

    # Conduct training
    for epoch in range(args.epochs):
        total_loss = 0

        if args.verbose:
            print('=============== Epoch {} ==============='.format(epoch+1))

        for it, (covers, secrets) in enumerate(train_loader):
            print('\tIter: {} / {}'.format(it+1, num_iters))

            # Enable CUDA if specified
            covers, secrets = covers.to(device), secrets.to(device)

            # Forward run
            prepped, container, revealed = model.forward(covers, secrets, device=device,
                                                         noise_level=args.noise_level)

            # Compute loss
            loss = criterion(container, covers)
            loss += args.beta * criterion(revealed, secrets)

            if args.gamma > 0:
                # Normalize each image in a minibatch of {covers, container}
                size = container.size()
                means = torch.tensor([0.485, 0.456, 0.406])
                stds = torch.tensor([0.229, 0.224, 0.225])
                means, stds = [item.view(3,1,1).expand(container.size()).to(device) for item in [means, stds]]
                normalized = [torch.div(torch.add(minibatch,-1,means),stds).to(device) \
                                for minibatch in [container, covers]]
                loss += args.gamma * criterion(*normalized)

            # Do back-propagation
            model.zero_grad()
            loss.backward()
            optimizer.step()

            # Collect statistics
            total_loss += loss.item()

            if args.verbose:
                print("\t\tLoss sum at iter {}: {}".format(it+1, loss.item()))

            if args.max_dataset_size > 0 and (it+1) * args.batch_size >= args.max_dataset_size:
                break

        if args.verbose:
            print('Loss at epoch {}: {}'.format(epoch+1, total_loss / num_iters))

        # Report statistics to NSML (if not draft)
        if not args.local:
            nsml.report(summary = True,
                        step = (epoch+1) * len(train_loader),
                        epoch_total = args.epochs,
                        train__loss = total_loss / num_iters)

            # Visualize input & output images
            for idx in range(args.viz_count):
                images = [images.detach()[idx] \
                            for images in [covers, secrets, container, revealed]]   # select example images
                visualize(viz, images, epoch,
                          c_mode=args.c_mode, s_mode = args.s_mode,
                          c_draw_rgb=args.c_draw_rgb, s_draw_rgb = args.s_draw_rgb,
                          use_cuda = use_cuda)

        # Save session if minimum loss is renewed
        if not args.draft and not args.local and total_loss < min_loss:
            nsml.save(epoch)
            min_loss = total_loss

        if args.draft:
            break
Ejemplo n.º 20
0
def train_distill(opts, train_loader, unlabel_loader, model, criterion,
                  optimizer, ema_optimizer, epoch, use_gpu, scheduler,
                  is_mixsim):
    global global_step
    scaler = torch.cuda.amp.GradScaler()

    losses = AverageMeter()
    losses_x = AverageMeter()
    losses_un = AverageMeter()

    losses_curr = AverageMeter()
    losses_x_curr = AverageMeter()
    losses_un_curr = AverageMeter()

    weight_scale = AverageMeter()
    acc_top1 = AverageMeter()
    acc_top5 = AverageMeter()

    model.train()

    # nCnt =0
    out = False
    local_step = 0
    while not out:
        labeled_train_iter = iter(train_loader)
        unlabeled_train_iter = iter(unlabel_loader)
        for batch_idx in range(len(train_loader)):
            try:
                data = labeled_train_iter.next()
                inputs_x, targets_x = data
            except:
                labeled_train_iter = iter(train_loader)
                data = labeled_train_iter.next()
                inputs_x, targets_x = data
            try:
                data = unlabeled_train_iter.next()
                inputs_u1, inputs_u2 = data
            except:
                unlabeled_train_iter = iter(unlabel_loader)
                data = unlabeled_train_iter.next()
                inputs_u1, inputs_u2 = data

            batch_size = inputs_x.size(0)
            # Transform label to one-hot
            classno = NUM_CLASSES
            targets_org = targets_x
            targets_x = torch.zeros(batch_size,
                                    classno).scatter_(1, targets_x.view(-1, 1),
                                                      1)

            if use_gpu:
                # Send input value to device 0, where first parameters of MixSim_Model is.
                if is_mixsim:
                    dev0 = 'cuda:{}'.format(opts.gpu_ids.split(',')[0])
                else:
                    dev0 = 'cuda'
                inputs_x, targets_x = inputs_x.to(dev0), targets_x.to(dev0)
                inputs_u1, inputs_u2 = inputs_u1.to(dev0), inputs_u2.to(dev0)

            with torch.no_grad():
                # compute guessed labels of unlabel samples
                _, pred_u1 = model(inputs_u1)
                _, pred_u2 = model(inputs_u2)
                pred_u_all = (torch.softmax(pred_u1, dim=1) +
                              torch.softmax(pred_u2, dim=1)) / 2
                pt = pred_u_all**(1 / opts.T)
                targets_u = pt / pt.sum(dim=1, keepdim=True)
                targets_u = targets_u.detach()

            # mixup
            all_inputs = torch.cat([inputs_x, inputs_u1, inputs_u2], dim=0)
            all_targets = torch.cat([targets_x, targets_u, targets_u], dim=0)

            lamda = np.random.beta(opts.alpha, opts.alpha)
            lamda = max(lamda, 1 - lamda)
            newidx = torch.randperm(all_inputs.size(0))
            input_a, input_b = all_inputs, all_inputs[newidx]
            target_a, target_b = all_targets, all_targets[newidx]

            mixed_input = lamda * input_a + (1 - lamda) * input_b
            mixed_target = lamda * target_a + (1 - lamda) * target_b

            # interleave labeled and unlabed samples between batches to get correct batchnorm calculation
            mixed_input = list(torch.split(mixed_input, batch_size))
            mixed_input = interleave(mixed_input, batch_size)

            optimizer.zero_grad()

            with torch.cuda.amp.autocast():
                _, logits_temp = model(mixed_input[0])
                logits = [logits_temp]
                for newinput in mixed_input[1:]:
                    _, logits_temp = model(newinput)
                    logits.append(logits_temp)

            # put interleaved samples back
            logits = interleave(logits, batch_size)
            logits_x = logits[0]
            logits_u = torch.cat(logits[1:], dim=0)

            loss_x, loss_un, weigts_mixing = criterion(
                logits_x, mixed_target[:batch_size], logits_u,
                mixed_target[batch_size:],
                epoch + batch_idx / len(train_loader), opts.epochs)
            loss = loss_x + weigts_mixing * loss_un

            losses.update(loss.item(), inputs_x.size(0))
            losses_x.update(loss_x.item(), inputs_x.size(0))
            losses_un.update(loss_un.item(), inputs_x.size(0))
            weight_scale.update(weigts_mixing, inputs_x.size(0))

            losses_curr.update(loss.item(), inputs_x.size(0))
            losses_x_curr.update(loss_x.item(), inputs_x.size(0))
            losses_un_curr.update(loss_un.item(), inputs_x.size(0))

            # compute gradient and do SGD step using amp
            scaler.scale(loss).backward()
            scaler.step(optimizer)
            ema_optimizer.step()
            scaler.update()
            scheduler.step()

            with torch.no_grad():
                # compute guessed labels of unlabel samples
                _, pred_x1 = model(inputs_x)

            if IS_ON_NSML and global_step % opts.log_interval == 0:
                nsml.report(step=global_step,
                            loss=losses_curr.avg,
                            loss_x=losses_x_curr.avg,
                            loss_un=losses_un_curr.avg)
                losses_curr.reset()
                losses_x_curr.reset()
                losses_un_curr.reset()

            acc_top1b = top_n_accuracy_score(targets_org.data.cpu().numpy(),
                                             pred_x1.data.cpu().numpy(),
                                             n=1) * 100
            acc_top5b = top_n_accuracy_score(targets_org.data.cpu().numpy(),
                                             pred_x1.data.cpu().numpy(),
                                             n=5) * 100
            acc_top1.update(torch.as_tensor(acc_top1b), inputs_x.size(0))
            acc_top5.update(torch.as_tensor(acc_top5b), inputs_x.size(0))

            local_step += 1
            global_step += 1

            if local_step >= opts.steps_per_epoch:
                out = True
                break

    return losses.avg, losses_x.avg, losses_un.avg, acc_top1.avg, acc_top5.avg
Ejemplo n.º 21
0
            train_loss = float(avg_loss / total_batch)
            val_loss = sess.run(textRNN.loss_val,
                                feed_dict={
                                    textRNN.input_y: y_test,
                                    textRNN.input_x: x_test,
                                    textRNN.dropout_keep_prob: 1
                                })
            if val_loss < min_val_loss:
                min_val_loss = val_loss

            print('epoch:', epoch, ' train_loss:', train_loss, ' val_loss:',
                  val_loss)

            nsml.report(summary=True,
                        scope=locals(),
                        epoch=epoch,
                        epoch_total=config.epochs,
                        train__loss=float(avg_loss / total_batch),
                        step=epoch)

            # DONOTCHANGE (You can decide how often you want to save the model)
            nsml.save(epoch)

    # 로컬 테스트 모드일때 사용합니다
    # 결과가 아래와 같이 나온다면, nsml submit을 통해서 제출할 수 있습니다.
    # [(0.3, 0), (0.7, 1), ... ]
    elif config.mode == 'test_local':
        with open(os.path.join(DATASET_PATH, 'train/train_data'),
                  'rt',
                  encoding='utf-8') as f:
            reviews = f.readlines()
        res = nsml.infer(reviews)
Ejemplo n.º 22
0
        data, y = data_loader(img_path, label_path)
        X = preprocessing(data)

        # Data loader
        batch_loader = DataLoader(
            dataset=MammoDataset(X, y),  ## pytorch data loader 사용
            batch_size=batch_size,
            shuffle=True)

        # Train the model
        total_step = len(batch_loader)
        for epoch in range(num_epochs):
            for i, (images, labels) in enumerate(batch_loader):
                images = images.to(device)
                labels = labels.to(device)

                # Forward pass
                outputs = model(images)
                loss = criterion(outputs, labels)

                # Backward and optimize
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()

            nsml.report(summary=True,
                        step=epoch,
                        epoch_total=num_epochs,
                        loss=loss.item())  #, acc=train_acc)
            nsml.save(epoch)
Ejemplo n.º 23
0
                if step % 150 == 0 or (config.debug and step % 1 == 0):
                    infer_result = local_infer(queries, references, queries_img, reference_img, batch_size)
                    mAP, mean_recall_at_K, min_first_1_at_K = evaluate_rank(infer_result)

                    if best_min_first_K >= min_first_1_at_K:
                        best_min_first_K = min_first_1_at_K
                        best_min_first_K_step = step
                        print("----> First_K @ 1 recall : %d / %d" % (min_first_1_at_K, len(mean_recall_at_K)))
                        do_save = True

                    if best_mAP <= mAP :
                        best_mAP = mAP
                        print("----> Best mAP : best-mAP {:g}".format(best_mAP))
                        do_save = True

                if epoch - prev_epoch == 1:
                    print("----> Epoch changed saving")
                    do_save = True

                if do_save:
                    # save model
                    nsml.report(summary=True, epoch=str(step), epoch_total=nb_epoch)
                    nsml.save(step)
                    print("Model saved : %d step" % step)
                    print("=============================================================================================================")


            except tf.errors.OutOfRangeError:
                print("finish train!")
                break
Ejemplo n.º 24
0
def trainTrans(model,
               total_batch_size,
               queue,
               criterion,
               optimizer,
               device,
               train_begin,
               train_loader_count,
               print_batch=5,
               teacher_forcing_ratio=1,
               visual=None,
               label_smoothing=0.1):
    total_loss = 0.
    total_num = 0
    total_dist = 0
    total_length = 0
    total_sent_num = 0
    batch = 0

    model.train()

    logger.info('train() start')

    begin = epoch_begin = time.time()
    while True:
        if queue.empty():
            logger.debug('queue is empty')

        feats, scripts, feat_lengths, script_lengths = queue.get()

        if feats.shape[0] == 0:
            # empty feats means closing one loader
            train_loader_count -= 1
            logger.debug('left train_loader: %d' % (train_loader_count))
            if train_loader_count == 0:
                break
            else:
                continue

        optimizer.zero_grad()

        feats = feats.to(device)
        scripts = scripts.to(device)

        src_len = scripts.size(1)
        target = scripts[:, 1:]

        pred, gold = model(feats, feat_lengths, scripts)
        loss, n_correct = cal_performance(pred,
                                          gold,
                                          smoothing=label_smoothing)
        y_hat = pred.max(-1)[1]

        total_loss += loss.item()
        total_num += sum(feat_lengths)

        display = random.randrange(0, 100) == 0
        dist, length = get_distance(target, y_hat, display=display)
        total_dist += dist
        total_length += length

        total_sent_num += target.size(0)

        loss.backward()
        optimizer.step()

        if visual:
            vis_log = {
                'Train Loss': total_loss / total_num,
                'Train CER': total_dist / total_length
            }
            visual.log(vis_log)

        if batch % print_batch == 0:
            current = time.time()
            elapsed = current - begin
            epoch_elapsed = (current - epoch_begin) / 60.0
            train_elapsed = (current - train_begin) / 3600.0

            logger.info(
                'batch: {:4d}/{:4d}, loss: {:.4f}, cer: {:.2f}, elapsed: {:.2f}s {:.2f}m {:.2f}h'
                .format(
                    batch,
                    #len(dataloader),
                    total_batch_size,
                    total_loss / total_num,
                    total_dist / total_length,
                    elapsed,
                    epoch_elapsed,
                    train_elapsed))
            begin = time.time()

            nsml.report(False,
                        step=trainTrans.cumulative_batch_count,
                        train_step__loss=total_loss / total_num,
                        train_step__cer=total_dist / total_length)
        batch += 1
        trainTrans.cumulative_batch_count += 1

    logger.info('train() completed')
    return total_loss / total_num, total_dist / total_length
Ejemplo n.º 25
0
    def run(self):
        mb_obs, mb_td_targets, mb_base_actions, \
        mb_xy0, mb_xy1, \
        mb_values, mb_dones \
          = [],[],[],[],[],[], []
        # ,[],[],[],[],[],[],[],[],[],[],[]

        mb_states = self.states
        for n in range(self.nsteps):
            # pi, pi2, x1, y1, x2, y2, v0
            pi1, pi_xy0, pi_xy1, values, states = self.model.step(
                self.obs, self.states, self.dones)

            pi1_noise = np.random.random_sample((self.nenv, 2)) * 0.3
            # avail = self.env.available_actions()
            # print("pi1 : ", pi1)
            # print("pi1 * self.base_act_mask : ", pi1 * self.base_act_mask)
            # print("pi1 * self.base_act_mask + pi1_noise : ", pi1 * self.base_act_mask + pi1_noise)

            base_actions = np.argmax(pi1 * self.base_act_mask + pi1_noise,
                                     axis=1)
            xy0 = np.argmax(pi_xy0, axis=1)

            x0 = (xy0 % 32).astype(int)
            y0 = (xy0 / 32).astype(int)

            xy1 = np.argmax(pi_xy1, axis=1)
            x1 = (xy1 % 32).astype(int)
            y1 = (xy1 / 32).astype(int)

            # pi (2?, 524) * (2?, 524) masking
            # print("base_actions : ", base_actions)
            # print("base_action_spec : ", base_action_spec)
            # sub1_act_mask, sub2_act_mask, sub3_act_mask = self.get_sub_act_mask(base_action_spec)
            # print("base_actions : ", base_actions, "base_action_spec", base_action_spec,
            #       "sub1_act_mask :", sub1_act_mask, "sub2_act_mask :", sub2_act_mask, "sub3_act_mask :", sub3_act_mask)
            # sub3_actions = np.argmax(pi_sub3, axis=1) # pi (2?, 2) [1 0]
            # sub4_actions = np.argmax(pi_sub4, axis=1) # pi (2?, 5) [4 4]
            # sub5_actions = np.argmax(pi_sub5, axis=1) # pi (2?, 10) [1 4]
            # sub6_actions = np.argmax(pi_sub6, axis=1) # pi (2?, 4) [3 1]
            # sub7_actions = np.argmax(pi_sub7, axis=1) # pi (2?, 2)
            # sub8_actions = np.argmax(pi_sub8, axis=1) # pi (2?, 4)
            # sub9_actions = np.argmax(pi_sub9, axis=1) # pi (2?, 500)
            # sub10_actions = np.argmax(pi_sub10, axis=1) # pi (2?, 4)
            # sub11_actions = np.argmax(pi_sub11, axis=1) # pi (2?, 10)
            # sub12_actions = np.argmax(pi_sub12, axis=1) # pi (2?, 500)

            # Scripted Agent Hacking

            for env_num in range(self.nenv):
                if (env_num >= self.nscripts):  # only for scripted agents
                    continue

                ob = self.obs[env_num, :, :, :]
                # extra = ob[:,:,-1]
                # selected = ob[:, :, -2]
                player_relative = ob[:, :, -1]

                #if(common.check_group_list())
                self.group_list[env_num] = common.update_group_list2(
                    self.control_groups[env_num])
                # if(len(self.action_queue[env_num]) == 0 and len(self.group_list[env_num]) == 0):
                #
                #   # Scripted Agent is only for even number agents
                #   self.action_queue[env_num] = common.group_init_queue(player_relative)

                if (len(self.action_queue[env_num]) == 0):

                    self.action_queue[env_num], self.group_id[env_num], self.dest_per_marine[env_num], self.xy_per_marine[env_num] = \
                      common.solve_tsp(player_relative,
                                       self.selected[env_num][0],
                                       self.group_list[env_num],
                                       self.group_id[env_num],
                                       self.dest_per_marine[env_num],
                                       self.xy_per_marine[env_num])

                base_actions[env_num] = 0
                x0[env_num] = 0
                y0[env_num] = 0
                x1[env_num] = 0
                y1[env_num] = 0

                if (len(self.action_queue[env_num]) > 0):
                    action = self.action_queue[env_num].pop(0)
                    # print("action :", action)
                    base_actions[env_num] = action.get("base_action", 0)

                    x0[env_num] = action.get("x0", 0)
                    y0[env_num] = action.get("y0", 0)
                    xy0[env_num] = y0[env_num] * 32 + x0[env_num]

                    x1[env_num] = action.get("x1", 0)
                    y1[env_num] = action.get("y1", 0)
                    xy1[env_num] = y1[env_num] * 32 + x1[env_num]

            base_actions = self.valid_base_action(base_actions)
            # print("valid_base_actions : ", base_actions)
            new_base_actions = self.trans_base_actions(base_actions)
            # print("new_base_actions : ", new_base_actions)

            base_action_spec = self.env.action_spec(new_base_actions)

            actions = self.construct_action(
                base_actions,
                base_action_spec,
                # sub3_actions, sub4_actions, sub5_actions,
                # sub6_actions,
                # sub7_actions, sub8_actions,
                # sub9_actions, sub10_actions,
                # sub11_actions, sub12_actions,
                x0,
                y0,
                x1,
                y1
                # , x2, y2
            )

            mb_obs.append(np.copy(self.obs))
            mb_base_actions.append(base_actions)
            # mb_sub3_actions.append(sub3_actions)
            # mb_sub4_actions.append(sub4_actions)
            # mb_sub5_actions.append(sub5_actions)
            # mb_sub6_actions.append(sub6_actions)
            # mb_sub7_actions.append(sub7_actions)
            # mb_sub8_actions.append(sub8_actions)
            # mb_sub9_actions.append(sub9_actions)
            # mb_sub10_actions.append(sub10_actions)
            # mb_sub11_actions.append(sub11_actions)
            # mb_sub12_actions.append(sub12_actions)

            mb_xy0.append(xy0)
            # mb_y0.append(y0)
            mb_xy1.append(xy1)
            # mb_y1.append(y1)
            # mb_x2.append(x2)
            # mb_y2.append(y2)
            mb_values.append(values)
            mb_dones.append(self.dones)

            #print("final acitons : ", actions)
            obs, rewards, dones, available_actions, army_counts, control_groups, selected, xy_per_marine = self.env.step(
                actions=actions)
            self.army_counts = army_counts
            self.control_groups = control_groups
            self.selected = selected
            for env_num, data in enumerate(xy_per_marine):
                # print("env_num", env_num, "xy_per_marine:", data)
                self.xy_per_marine[env_num] = data
            self.update_available(available_actions)

            self.states = states
            self.dones = dones
            for n, done in enumerate(dones):
                self.total_reward[n] += float(rewards[n])
                if done:
                    self.obs[n] = self.obs[n] * 0
                    self.episodes += 1
                    num_episodes = self.episodes
                    self.episode_rewards.append(self.total_reward[n])

                    mean_100ep_reward = round(
                        np.mean(self.episode_rewards[-101:-1]), 1)

                    if (n < self.nscripts):  # scripted agents
                        self.episode_rewards_script.append(
                            self.total_reward[n])
                        mean_100ep_reward_script = round(
                            np.mean(self.episode_rewards_script[-101:-1]), 1)
                        # logger.record_tabular("reward script",
                        #                       self.total_reward[n])
                        # logger.record_tabular("mean reward script",
                        #                       mean_100ep_reward_script)
                        nsml.report(
                            reward_script=self.total_reward[n],
                            mean_reward_script=mean_100ep_reward_script,
                            reward=self.total_reward[n],
                            mean_100ep_reward=mean_100ep_reward,
                            episodes=self.episodes,
                            step=self.episodes,
                            scope=locals())
                    else:
                        self.episode_rewards_a2c.append(self.total_reward[n])
                        mean_100ep_reward_a2c = round(
                            np.mean(self.episode_rewards_a2c[-101:-1]), 1)
                        # logger.record_tabular("reward a2c",
                        #                       self.total_reward[n])
                        # logger.record_tabular("mean reward a2c",
                        #                       mean_100ep_reward_a2c)
                        nsml.report(reward_a2c=self.total_reward[n],
                                    mean_reward_a2c=mean_100ep_reward_a2c,
                                    reward=self.total_reward[n],
                                    mean_100ep_reward=mean_100ep_reward,
                                    episodes=self.episodes,
                                    step=self.episodes,
                                    scope=locals())

                    #print("env %s done! reward : %s mean_100ep_reward : %s " %
                    #      (n, self.total_reward[n], mean_100ep_reward))

                    # logger.record_tabular("reward", self.total_reward[n])
                    # logger.record_tabular("mean 100 episode reward",
                    #                       mean_100ep_reward)
                    # logger.record_tabular("episodes", self.episodes)

                    # logger.dump_tabular()

                    self.total_reward[n] = 0
                    self.group_list[n] = []

                    model = self.model
                    if self.callback is not None:
                        self.callback(locals(), globals())

            #print("rewards : ", rewards)
            #print("self.total_reward :", self.total_reward)
            self.update_obs(obs)
            mb_td_targets.append(rewards)
        mb_dones.append(self.dones)
        #batch of steps to batch of rollouts
        mb_obs = np.asarray(mb_obs, dtype=np.uint8).swapaxes(1, 0).reshape(
            self.batch_ob_shape)
        mb_td_targets = np.asarray(mb_td_targets,
                                   dtype=np.float32).swapaxes(1, 0)
        mb_base_actions = np.asarray(mb_base_actions,
                                     dtype=np.int32).swapaxes(1, 0)
        # mb_sub3_actions = np.asarray(mb_sub3_actions, dtype=np.int32).swapaxes(1, 0)
        # mb_sub4_actions = np.asarray(mb_sub4_actions, dtype=np.int32).swapaxes(1, 0)
        # mb_sub5_actions = np.asarray(mb_sub5_actions, dtype=np.int32).swapaxes(1, 0)
        # mb_sub6_actions = np.asarray(mb_sub6_actions, dtype=np.int32).swapaxes(1, 0)
        # mb_sub7_actions = np.asarray(mb_sub7_actions, dtype=np.int32).swapaxes(1, 0)
        # mb_sub8_actions = np.asarray(mb_sub8_actions, dtype=np.int32).swapaxes(1, 0)
        # mb_sub9_actions = np.asarray(mb_sub9_actions, dtype=np.int32).swapaxes(1, 0)
        # mb_sub10_actions = np.asarray(mb_sub10_actions, dtype=np.int32).swapaxes(1, 0)
        # mb_sub11_actions = np.asarray(mb_sub11_actions, dtype=np.int32).swapaxes(1, 0)
        # mb_sub12_actions = np.asarray(mb_sub12_actions, dtype=np.int32).swapaxes(1, 0)

        mb_xy0 = np.asarray(mb_xy0, dtype=np.int32).swapaxes(1, 0)
        # mb_y0 = np.asarray(mb_y0, dtype=np.int32).swapaxes(1, 0)
        mb_xy1 = np.asarray(mb_xy1, dtype=np.int32).swapaxes(1, 0)
        # mb_y1 = np.asarray(mb_y1, dtype=np.int32).swapaxes(1, 0)
        # mb_x2 = np.asarray(mb_x2, dtype=np.int32).swapaxes(1, 0)
        # mb_y2 = np.asarray(mb_y2, dtype=np.int32).swapaxes(1, 0)

        mb_values = np.asarray(mb_values, dtype=np.float32).swapaxes(1, 0)
        mb_dones = np.asarray(mb_dones, dtype=np.bool).swapaxes(1, 0)
        mb_masks = mb_dones[:, :-1]
        mb_dones = mb_dones[:, 1:]
        last_values = self.model.value(self.obs, self.states,
                                       self.dones).tolist()
        #discount/bootstrap off value fn
        for n, (rewards, dones,
                value) in enumerate(zip(mb_td_targets, mb_dones, last_values)):
            rewards = rewards.tolist()
            dones = dones.tolist()
            if dones[-1] == 0:
                rewards = discount_with_dones(rewards + [value], dones + [0],
                                              self.gamma)[:-1]
            else:
                rewards = discount_with_dones(rewards, dones, self.gamma)
            mb_td_targets[n] = rewards
        mb_td_targets = mb_td_targets.flatten()
        mb_base_actions = mb_base_actions.flatten()
        # mb_sub3_actions = mb_sub3_actions.flatten()
        # mb_sub4_actions = mb_sub4_actions.flatten()
        # mb_sub5_actions = mb_sub5_actions.flatten()
        # mb_sub6_actions = mb_sub6_actions.flatten()
        # mb_sub7_actions = mb_sub7_actions.flatten()
        # mb_sub8_actions = mb_sub8_actions.flatten()
        # mb_sub9_actions = mb_sub9_actions.flatten()
        # mb_sub10_actions = mb_sub10_actions.flatten()
        # mb_sub11_actions = mb_sub11_actions.flatten()
        # mb_sub12_actions = mb_sub12_actions.flatten()
        mb_xy0 = mb_xy0.flatten()
        # mb_y0 = mb_y0.flatten()
        mb_xy1 = mb_xy1.flatten()
        # mb_y1 = mb_y1.flatten()
        # mb_x2 = mb_x2.flatten()
        # mb_y2 = mb_y2.flatten()

        mb_values = mb_values.flatten()
        mb_masks = mb_masks.flatten()
        return mb_obs, mb_states, mb_td_targets, mb_masks, \
               mb_base_actions, mb_xy0, mb_xy1, mb_values
Ejemplo n.º 26
0
            for iter_, (image, image_id, label) in enumerate(train_loader):
                image = image.cuda()
                label = label.cuda()

                pred = model(image)
                loss = criterion(pred, label)

                total_loss += loss.item() * image.size(0)
                num_images += image.size(0)

                optimizer.zero_grad()
                loss.backward()
                optimizer.step()

            loss_average = total_loss / float(num_images)
            scheduler.step(metrics=loss_average, epoch=epoch)

            if IS_ON_NSML:
                nsml.save(str(epoch + 1))

            gt_label = os.path.join(DATASET_PATH, 'train/train_data/val_label')
            acc = local_eval(model, val_loader, gt_label)
            print(f'[{epoch + 1}/{config.num_epochs}] '
                  f'Validation performance: {acc:.3f}')
            nsml.report(step=epoch, val_acc=acc)
            nsml.report(step=epoch,
                        train_loss=loss_average,
                        learning_rate=optimizer.param_groups[0]['lr'])
            if IS_ON_NSML and acc > best_accuracy:
                nsml.save(str(epoch + 1))
                best_accuracy = acc
Ejemplo n.º 27
0
def learn(policy,
          env,
          seed,
          total_timesteps=int(40e6),
          gamma=0.99,
          log_interval=1,
          nprocs=24,
          nscripts=12,
          nsteps=20,
          nstack=4,
          ent_coef=0.01,
          vf_coef=0.5,
          vf_fisher_coef=1.0,
          lr=0.25,
          max_grad_norm=0.01,
          kfac_clip=0.001,
          save_interval=None,
          lrschedule='linear',
          callback=None):
    tf.reset_default_graph()
    set_global_seeds(seed)

    nenvs = nprocs
    ob_space = (32, 32, 3)  # env.observation_space
    ac_space = (32, 32)
    make_model = lambda: Model(policy,
                               ob_space,
                               ac_space,
                               nenvs,
                               total_timesteps,
                               nprocs=nprocs,
                               nscripts=nscripts,
                               nsteps=nsteps,
                               nstack=nstack,
                               ent_coef=ent_coef,
                               vf_coef=vf_coef,
                               vf_fisher_coef=vf_fisher_coef,
                               lr=lr,
                               max_grad_norm=max_grad_norm,
                               kfac_clip=kfac_clip,
                               lrschedule=lrschedule)

    if save_interval and logger.get_dir():
        import cloudpickle
        with open(osp.join(logger.get_dir(), 'make_model.pkl'), 'wb') as fh:
            fh.write(cloudpickle.dumps(make_model))
    model = make_model()
    print("make_model complete!")
    runner = Runner(env,
                    model,
                    nsteps=nsteps,
                    nscripts=nscripts,
                    nstack=nstack,
                    gamma=gamma,
                    callback=callback)
    nbatch = nenvs * nsteps
    tstart = time.time()
    #enqueue_threads = model.q_runner.create_threads(model.sess, coord=tf.train.Coordinator(), start=True)
    for update in range(1, total_timesteps // nbatch + 1):
        # print('iteration: %d' % update)
        obs, states, td_targets, masks, actions, xy0, xy1, values = runner.run(
        )
        # print('values: %s' % values)

        policy_loss, value_loss, policy_entropy, \
        policy_loss_xy0, policy_entropy_xy0, \
        policy_loss_xy1, policy_entropy_xy1, \
          = model.train(obs, states, td_targets,
                        masks, actions,
                        xy0, xy1, values)

        model.old_obs = obs
        nseconds = time.time() - tstart
        fps = int((update * nbatch) / nseconds)
        if update % log_interval == 0 or update == 1:
            ev = explained_variance(values, td_targets)
            nsml.report(nupdates=update,
                        total_timesteps=update * nbatch,
                        fps=fps,
                        policy_entropy=float(policy_entropy),
                        policy_loss=float(policy_loss),
                        policy_loss_xy0=float(policy_loss_xy0),
                        policy_entropy_xy0=float(policy_entropy_xy0),
                        policy_loss_xy1=float(policy_loss_xy1),
                        policy_entropy_xy1=float(policy_entropy_xy1),
                        value_loss=float(value_loss),
                        explained_variance=float(ev),
                        batch_size=nbatch,
                        step=update * nbatch,
                        scope=locals())
            # logger.record_tabular("nupdates", update)
            # logger.record_tabular("total_timesteps", update * nbatch)
            # logger.record_tabular("fps", fps)
            # logger.record_tabular("policy_entropy", float(policy_entropy))
            # logger.record_tabular("policy_loss", float(policy_loss))

            # logger.record_tabular("policy_loss_xy0", float(policy_loss_xy0))
            # logger.record_tabular("policy_entropy_xy0",
            #                       float(policy_entropy_xy0))
            # logger.record_tabular("policy_loss_xy1", float(policy_loss_xy1))
            # logger.record_tabular("policy_entropy_xy1",
            #                       float(policy_entropy_xy1))
            # # logger.record_tabular("policy_loss_y0", float(policy_loss_y0))
            # # logger.record_tabular("policy_entropy_y0", float(policy_entropy_y0))

            # logger.record_tabular("value_loss", float(value_loss))
            # logger.record_tabular("explained_variance", float(ev))
            # logger.dump_tabular()

        if save_interval and (update % save_interval == 0
                              or update == 1) and logger.get_dir():
            savepath = osp.join(logger.get_dir(), 'checkpoint%.5i' % update)
            print('Saving to', savepath)
            model.save(savepath)

    env.close()
Ejemplo n.º 28
0
def main():
    global char2index
    global index2char
    global SOS_token
    global EOS_token
    global PAD_token

    parser = argparse.ArgumentParser(description='Speech hackathon lilililill model')
    parser.add_argument('--max_epochs', type=int, default=1000, help='number of max epochs in training (default: 1000)')
    parser.add_argument('--no_cuda', action='store_true', default=False, help='disables CUDA training')
    parser.add_argument('--save_name', type=str, default='model', help='the name of model in nsml or local')

    # parser.add_argument('--dropout', type=float, default=0.2, help='dropout rate in training (default: 0.2)')
    parser.add_argument('--lr_1', type=float, default=1e-04, help='learning rate 1 (default: 0.0001)')
    parser.add_argument('--lr_2', type=float, default=1e-04, help='learning rate 2 (default: 0.0001)')
    parser.add_argument('--num_mels', type=int, default=160, help='number of the mel bands (default: 160)')
    parser.add_argument('--num_layers', type=int, default=1, help='number of layers (default: 1)')
    parser.add_argument('--batch_size', type=int, default=128, help='batch size in training (default: 128)')
    parser.add_argument("--num_thread", type=int, default=4, help='number of the loading thread (default: 4)')
    parser.add_argument('--num_hidden_enc', type=int, default=1024, help='hidden size of model (default: 1024)')
    parser.add_argument('--num_hidden_dec', type=int, default=512, help='hidden size of model decoder (default: 512)')
    parser.add_argument('--num_hidden_seq', type=int, default=1024, help='hidden size of model seq2seq (default: 1024)')
    parser.add_argument('--nsc_in_ms', type=int, default=40,
                        help='Number of sample size per time segment in ms (default: 40)')

    parser.add_argument('--ref_repeat', type=int, default=1,
                        help='Number of repetition of reference seq2seq (default: 1)')
    parser.add_argument('--loss_lim', type=float, default=0.05, help='Minimum loss threshold (default: 0.05)')

    parser.add_argument('--mode', type=str, default='train')
    parser.add_argument("--pause", type=int, default=0)
    parser.add_argument('--memo', type=str, default='', help='Comment you wish to leave')

    parser.add_argument('--train_ratio', type=float, default=0.5, help='Training ratio (default: 0.5)')

    parser.add_argument('--load', type=str, default=None)

    parser.add_argument('--debug', type=str, default='False', help='debug mode')

    args = parser.parse_args()

    batch_size = args.batch_size
    num_thread = args.num_thread
    num_mels = args.num_mels

    char2index, index2char = load_label('./hackathon.labels')
    SOS_token = char2index['<s>']  # '<sos>' or '<s>'
    EOS_token = char2index['</s>']  # '<eos>' or '</s>'
    PAD_token = char2index['_']  # '-' or '_'

    unicode_jamo_list = My_Unicode_Jamo_v2()

    tokenizer = Tokenizer(unicode_jamo_list)
    jamo_tokens = tokenizer.word2num(unicode_jamo_list)

    args.cuda = not args.no_cuda and torch.cuda.is_available()
    device = torch.device('cuda' if args.cuda else 'cpu')

    net = Mel2SeqNet_General(num_mels, args.num_hidden_enc, args.num_hidden_dec, len(unicode_jamo_list), args.num_layers, device)
    net_optimizer = optim.Adam(net.parameters(), lr=args.lr_1)
    ctc_loss = nn.CTCLoss().to(device)

    net_B = Seq2SeqNet_v2(args.num_hidden_seq, jamo_tokens, char2index, device)
    net_B_optimizer = optim.Adam(net_B.parameters(), lr=args.lr_2)
    net_B_criterion = nn.NLLLoss(reduction='none').to(device)

    bind_model(net, net_B, net_optimizer, net_B_optimizer, index2char, tokenizer)

    if args.pause == 1:
        nsml.paused(scope=locals())

    if args.mode != "train":
        return

    if args.load != None:
        nsml.load(checkpoint='model', session='team47/sr-hack-2019-50000/' + args.load)
        nsml.save('saved')

    for g in net_optimizer.param_groups:
        g['lr'] = args.lr_1
        logger.info('Learning rate of the net: {}'.format(g['lr']))

    for g in net_B_optimizer.param_groups:
        g['lr'] = args.lr_2
        logger.info('Learning rate of the net B: {}'.format(g['lr']))

    print(net)
    print(net_B)

    wav_paths, script_paths, korean_script_paths = get_paths(DATASET_PATH)

    logger.info('wav_paths len: {}'.format(len(wav_paths)))
    logger.info('script_paths len: {}'.format(len(script_paths)))
    logger.info('korean_script_paths len: {}'.format(len(korean_script_paths)))

    korean_script_list, jamo_script_list = get_korean_and_jamo_list_v2(korean_script_paths)

    logger.info('Korean script 0: {}'.format(korean_script_list[0]))
    logger.info('Korean script 0 length: {}'.format(len(korean_script_list[0])))
    logger.info('Jamo script 0: {}'.format(jamo_script_list[0]))
    logger.info('Jamo script 0 length: {}'.format(len(jamo_script_list[0])))

    script_path_list = get_script_list(script_paths, SOS_token, EOS_token)

    ground_truth_list = [(tokenizer.word2num(['<s>'] + list(jamo_script_list[i]) + ['</s>'])) for i in
                         range(len(jamo_script_list))]

    logger.info('Train Ratio: {}'.format(args.train_ratio))
    split_index = int(args.train_ratio * len(wav_paths))
    # split_index = int(0.9 * len(wav_paths))

    wav_path_list_train = wav_paths[:split_index]
    ground_truth_list_train = ground_truth_list[:split_index]
    korean_script_list_train = korean_script_list[:split_index]
    script_path_list_train = script_path_list[:split_index]

    wav_path_list_eval = wav_paths[split_index:]
    ground_truth_list_eval = ground_truth_list[split_index:]
    korean_script_list_eval = korean_script_list[split_index:]
    script_path_list_eval = script_path_list[split_index:]

    logger.info('Total:Train:Eval = {}:{}:{}'.format(len(wav_paths), len(wav_path_list_train), len(wav_path_list_eval)))

    preloader_train = Threading_Batched_Preloader_v2(wav_path_list_train, ground_truth_list_train,
                                                     script_path_list_train, korean_script_list_train, batch_size,
                                                     num_mels, args.nsc_in_ms, is_train=True)

    preloader_eval = Threading_Batched_Preloader_v2(wav_path_list_eval, ground_truth_list_eval, script_path_list_eval,
                                                    korean_script_list_eval, batch_size, num_mels, args.nsc_in_ms,
                                                    is_train=False)

    best_loss = 1e10
    best_eval_cer = 1e10

    # load all target scripts for reducing disk i/o
    target_path = os.path.join(DATASET_PATH, 'train_label')
    load_targets(target_path)

    logger.info('start')

    train_begin = time.time()

    for epoch in range(args.max_epochs):

        logger.info((datetime.now().strftime('%m-%d %H:%M:%S')))

        net.train()
        net_B.train()

        preloader_train.initialize_batch(num_thread)
        loss_list_train = list()
        seq2seq_loss_list_train = list()
        seq2seq_loss_list_train_ref = list()

        logger.info("Initialized Training Preloader")
        count = 0

        total_dist = 0
        total_length = 1
        total_dist_ref = 0
        total_length_ref = 1

        while not preloader_train.end_flag:
            batch = preloader_train.get_batch()
            if batch is not None:
                tensor_input, ground_truth, loss_mask, length_list, batched_num_script, batched_num_script_loss_mask = batch
                pred_tensor, loss = train(net, net_optimizer, ctc_loss, tensor_input.to(device),
                                          ground_truth.to(device), length_list.to(device), device)
                loss_list_train.append(loss)
                jamo_result = Decode_Prediction_No_Filtering(pred_tensor, tokenizer)
                true_string_list = Decode_Num_Script(batched_num_script.detach().cpu().numpy(), index2char)

                for i in range(args.ref_repeat):
                    lev_input_ref = ground_truth

                    lev_pred_ref, attentions_ref, seq2seq_loss_ref = net_B.net_train(lev_input_ref.to(device),
                                                                                     batched_num_script.to(device),
                                                                                     batched_num_script_loss_mask.to(
                                                                                         device),
                                                                                     net_B_optimizer,
                                                                                     net_B_criterion)

                    pred_string_list_ref = Decode_Lev_Prediction(lev_pred_ref, index2char)
                    seq2seq_loss_list_train_ref.append(seq2seq_loss_ref)
                    dist_ref, length_ref = char_distance_list(true_string_list, pred_string_list_ref)

                    pred_string_list = [None]

                    dist = 0
                    length = 0

                    if (loss < args.loss_lim):
                        lev_input = Decode_CTC_Prediction_And_Batch(pred_tensor)
                        lev_pred, attentions, seq2seq_loss = net_B.net_train(lev_input.to(device),
                                                                             batched_num_script.to(device),
                                                                             batched_num_script_loss_mask.to(device),
                                                                             net_B_optimizer, net_B_criterion)
                        pred_string_list = Decode_Lev_Prediction(lev_pred, index2char)
                        seq2seq_loss_list_train.append(seq2seq_loss)
                        dist, length = char_distance_list(true_string_list, pred_string_list)

                total_dist_ref += dist_ref
                total_length_ref += length_ref

                total_dist += dist
                total_length += length

                count += 1

                if count % 25 == 0:
                    logger.info(
                        "Train: Count {} | {} => {}".format(count, true_string_list[0], pred_string_list_ref[0]))

                    logger.info("Train: Count {} | {} => {} => {}".format(count, true_string_list[0], jamo_result[0],
                                                                          pred_string_list[0]))

            else:
                logger.info("Training Batch is None")

        train_loss = np.mean(np.asarray(loss_list_train))
        train_cer = np.mean(np.asarray(total_dist / total_length))
        train_cer_ref = np.mean(np.asarray(total_dist_ref / total_length_ref))

        logger.info("Mean Train Loss: {}".format(train_loss))
        logger.info("Total Train CER: {}".format(train_cer))
        logger.info("Total Train Reference CER: {}".format(train_cer_ref))

        preloader_eval.initialize_batch(num_thread)
        loss_list_eval = list()
        seq2seq_loss_list_eval = list()
        seq2seq_loss_list_eval_ref = list()

        logger.info("Initialized Evaluation Preloader")

        count = 0
        total_dist = 0
        total_length = 1
        total_dist_ref = 0
        total_length_ref = 1

        net.eval()
        net_B.eval()

        while not preloader_eval.end_flag:
            batch = preloader_eval.get_batch()
            if batch is not None:
                tensor_input, ground_truth, loss_mask, length_list, batched_num_script, batched_num_script_loss_mask = batch
                pred_tensor, loss = evaluate(net, ctc_loss, tensor_input.to(device), ground_truth.to(device),
                                             length_list.to(device), device)
                loss_list_eval.append(loss)

                jamo_result = Decode_Prediction_No_Filtering(pred_tensor, tokenizer)

                true_string_list = Decode_Num_Script(batched_num_script.detach().cpu().numpy(), index2char)

                lev_input_ref = ground_truth
                lev_pred_ref, attentions_ref, seq2seq_loss_ref = net_B.net_eval(lev_input_ref.to(device),
                                                                                batched_num_script.to(device),
                                                                                batched_num_script_loss_mask.to(device),
                                                                                net_B_criterion)

                pred_string_list_ref = Decode_Lev_Prediction(lev_pred_ref, index2char)
                seq2seq_loss_list_train_ref.append(seq2seq_loss_ref)
                dist_ref, length_ref = char_distance_list(true_string_list, pred_string_list_ref)

                pred_string_list = [None]

                dist = 0
                length = 0

                if (loss < args.loss_lim):
                    lev_input = Decode_CTC_Prediction_And_Batch(pred_tensor)
                    lev_pred, attentions, seq2seq_loss = net_B.net_eval(lev_input.to(device),
                                                                        batched_num_script.to(device),
                                                                        batched_num_script_loss_mask.to(device),
                                                                        net_B_criterion)
                    pred_string_list = Decode_Lev_Prediction(lev_pred, index2char)
                    seq2seq_loss_list_train.append(seq2seq_loss)
                    dist, length = char_distance_list(true_string_list, pred_string_list)

                total_dist_ref += dist_ref
                total_length_ref += length_ref

                total_dist += dist
                total_length += length

                count += 1

                if count % 10 == 0:
                    logger.info("Eval: Count {} | {} => {}".format(count, true_string_list[0], pred_string_list_ref[0]))

                    logger.info("Eval: Count {} | {} => {} => {}".format(count, true_string_list[0], jamo_result[0],
                                                                         pred_string_list[0]))

            else:
                logger.info("Training Batch is None")

        eval_cer = total_dist / total_length
        eval_cer_ref = total_dist_ref / total_length_ref
        eval_loss = np.mean(np.asarray(loss_list_eval))

        logger.info("Mean Evaluation Loss: {}".format(eval_loss))
        logger.info("Total Evaluation CER: {}".format(eval_cer))
        logger.info("Total Evaluation Reference CER: {}".format(eval_cer_ref))

        nsml.report(False, step=epoch, train_epoch__loss=train_loss, train_epoch__cer=train_cer,
                    train_epoch__cer_ref=train_cer_ref,
                    eval__loss=eval_loss, eval__cer=eval_cer, eval__cer_ref=eval_cer_ref)

        nsml.save(args.save_name)
        best_model = (eval_cer < best_eval_cer)
        if best_model:
            nsml.save('best')
            best_eval_cer = eval_cer
Ejemplo n.º 29
0
  def run(self):
    mb_obs, mb_td_targets, mb_base_actions, \
    mb_xy0, mb_xy1, \
    mb_values, mb_dones \
      = [], [], [], [], [], [], []

    mb_states = self.states
    for n in range(self.nsteps):
      # pi, pi2, x1, y1, x2, y2, v0
      pi1, pi_xy0, pi_xy1, values, states = self.model.step(
          self.obs, self.states, self.dones)

      pi1_noise = np.random.random_sample((self.nenv, 3)) * 0.3

      base_actions = np.argmax(
          pi1 * self.base_act_mask + pi1_noise, axis=1)
      xy0 = np.argmax(pi_xy0, axis=1)

      x0 = (xy0 % 32).astype(int)
      y0 = (xy0 / 32).astype(int)

      xy1 = np.argmax(pi_xy1, axis=1)
      x1 = (xy1 % 32).astype(int)
      y1 = (xy1 / 32).astype(int)

      # Scripted Agent Hacking

      for env_num in range(self.nenv):
        if env_num >= self.nscripts:  # only for scripted agents
          continue

        ob = self.obs[env_num, :, :, :]
        player_relative = ob[:, :, -1]

        self.group_list[env_num] = common.update_group_list2(
            self.control_groups[env_num])

        if len(self.action_queue[env_num]) == 0:

          self.action_queue[env_num], self.group_id[env_num], self.dest_per_marine[env_num], self.xy_per_marine[env_num] = \
            common.solve_tsp(player_relative,
                             self.selected[env_num][0],
                             self.group_list[env_num],
                             self.group_id[env_num],
                             self.dest_per_marine[env_num],
                             self.xy_per_marine[env_num])

        base_actions[env_num] = 0
        x0[env_num] = 0
        y0[env_num] = 0
        x1[env_num] = 0
        y1[env_num] = 0

        if len(self.action_queue[env_num]) > 0:
          action = self.action_queue[env_num].pop(0)
          base_actions[env_num] = action.get("base_action", 0)

          x0[env_num] = action.get("x0", 0)
          y0[env_num] = action.get("y0", 0)
          xy0[env_num] = y0[env_num] * 32 + x0[env_num]

          x1[env_num] = action.get("x1", 0)
          y1[env_num] = action.get("y1", 0)
          xy1[env_num] = y1[env_num] * 32 + x1[env_num]

      base_actions = self.valid_base_action(base_actions)
      new_base_actions = self.trans_base_actions(base_actions)

      base_action_spec = self.env.action_spec(new_base_actions)
      # print("base_actions:", base_actions)
      actions = self.construct_action(
          base_actions,
          base_action_spec,
          x0,
          y0,
          x1,
          y1
      )

      mb_obs.append(np.copy(self.obs))
      mb_base_actions.append(base_actions)

      mb_xy0.append(xy0)
      mb_xy1.append(xy1)
      mb_values.append(values)
      mb_dones.append(self.dones)

      #print("final acitons : ", actions)
      obs, rewards, dones,\
      available_actions, army_counts,\
      control_groups, selected, xy_per_marine\
      = self.env.step(
          actions=actions)
      self.army_counts = army_counts
      self.control_groups = control_groups
      self.selected = selected
      for env_num, data in enumerate(xy_per_marine):
        self.xy_per_marine[env_num] = data
      self.update_available(available_actions)

      self.states = states
      self.dones = dones
      mean_100ep_reward_a2c = 0
      for n, done in enumerate(dones):
        self.total_reward[n] += float(rewards[n])
        if done:
          self.obs[n] = self.obs[n] * 0
          self.episodes += 1
          num_episodes = self.episodes
          self.episode_rewards.append(self.total_reward[n])

          model = self.model
          mean_100ep_reward = round(
              np.mean(self.episode_rewards[-101:]), 1)
          if (n < self.nscripts):  # scripted agents
            self.episode_rewards_script.append(
                self.total_reward[n])
            mean_100ep_reward_script = round(
                np.mean(self.episode_rewards_script[-101:]), 1)
            nsml.report(
                reward_script=self.total_reward[n],
                mean_reward_script=mean_100ep_reward_script,
                reward=self.total_reward[n],
                mean_100ep_reward=mean_100ep_reward,
                episodes=self.episodes,
                step=self.episodes,
                scope=locals()
            )
          else:
            self.episode_rewards_a2c.append(self.total_reward[n])
            mean_100ep_reward_a2c = round(
                np.mean(self.episode_rewards_a2c[-101:]), 1)
            nsml.report(
                reward_a2c=self.total_reward[n],
                mean_reward_a2c=mean_100ep_reward_a2c,
                reward=self.total_reward[n],
                mean_100ep_reward=mean_100ep_reward,
                episodes=self.episodes,
                step=self.episodes,
                scope=locals()
            )
            print("mean_100ep_reward_a2c", mean_100ep_reward_a2c)

          if self.callback is not None:
            self.callback(locals(), globals())
          self.total_reward[n] = 0
          self.group_list[n] = []


      self.update_obs(obs)
      mb_td_targets.append(rewards)
    mb_dones.append(self.dones)
    #batch of steps to batch of rollouts
    mb_obs = np.asarray(
        mb_obs, dtype=np.uint8).swapaxes(1, 0).reshape(
            self.batch_ob_shape)
    mb_td_targets = np.asarray(mb_td_targets, dtype=np.float32).swapaxes(1, 0)
    mb_base_actions = np.asarray(
        mb_base_actions, dtype=np.int32).swapaxes(1, 0)

    mb_xy0 = np.asarray(mb_xy0, dtype=np.int32).swapaxes(1, 0)
    mb_xy1 = np.asarray(mb_xy1, dtype=np.int32).swapaxes(1, 0)

    mb_values = np.asarray(mb_values, dtype=np.float32).swapaxes(1, 0)
    mb_dones = np.asarray(mb_dones, dtype=np.bool).swapaxes(1, 0)
    mb_masks = mb_dones[:, :-1]
    mb_dones = mb_dones[:, 1:]
    last_values = self.model.value(self.obs, self.states,
                                   self.dones).tolist()
    #discount/bootstrap off value fn
    for n, (rewards, dones, value) in enumerate(
        zip(mb_td_targets, mb_dones, last_values)):
      rewards = rewards.tolist()
      dones = dones.tolist()
      if dones[-1] == 0:
        rewards = discount_with_dones(rewards + [value], dones + [0],
                                      self.gamma)[:-1]
      else:
        rewards = discount_with_dones(rewards, dones, self.gamma)
      mb_td_targets[n] = rewards
    mb_td_targets = mb_td_targets.flatten()
    mb_base_actions = mb_base_actions.flatten()
    mb_xy0 = mb_xy0.flatten()
    mb_xy1 = mb_xy1.flatten()

    mb_values = mb_values.flatten()
    mb_masks = mb_masks.flatten()
    return mb_obs, mb_states, mb_td_targets, mb_masks, \
           mb_base_actions, mb_xy0, mb_xy1, mb_values
Ejemplo n.º 30
0
def main():
    seed_everything()

    # yml = 'configs/base.yml'
    # config = utils.config.load(yml)
    # pprint.pprint(config, indent=2)

    model = get_model(config).cuda()
    bind_model(model)

    args = get_args()
    if args.pause:  ## test mode일 때
        print('Inferring Start...')
        nsml.paused(scope=locals())

    if args.mode == 'train':  ### training mode일 때
        print('Training Start...')

        # no bias decay
        if config.OPTIMIZER.NO_BIAS_DECAY:
            group_decay, group_no_decay = group_weight(model)
            params = [{'params': group_decay},
                      {'params': group_no_decay, 'weight_decay': 0.0}]
        else:
            params = model.parameters()

        optimizer = get_optimizer(config, params)
        optimizer.param_groups[0]['initial_lr'] = config.OPTIMIZER.LR
        if config.OPTIMIZER.NO_BIAS_DECAY:
            optimizer.param_groups[1]['initial_lr'] = config.OPTIMIZER.LR

        ###############################################################################################

        if IS_LOCAL:
            prepare_train_directories(config)
            utils.config.save_config(yml, config.LOCAL_TRAIN_DIR)

            checkpoint = utils.checkpoint.get_initial_checkpoint(config)
            if checkpoint is not None:
                last_epoch, score, loss = utils.checkpoint.load_checkpoint(config, model, checkpoint)
            else:
                print('[*] no checkpoint found')
                last_epoch, score, loss = -1, -1, float('inf')
            print('last epoch:{} score:{:.4f} loss:{:.4f}'.format(last_epoch, score, loss))

        else:
            last_epoch = -1
        ###############################################################################################

        scheduler = get_scheduler(config, optimizer, last_epoch=last_epoch)

        ###############################################################################################
        if IS_LOCAL:
            if config.SCHEDULER.NAME == 'multi_step':
                if config.SCHEDULER.WARMUP:
                    scheduler_dict = scheduler.state_dict()['after_scheduler'].state_dict()
                else:
                    scheduler_dict = scheduler.state_dict()

                milestones = scheduler_dict['milestones']
                step_count = len([i for i in milestones if i < last_epoch])
                optimizer.param_groups[0]['lr'] *= scheduler_dict['gamma'] ** step_count
                if config.OPTIMIZER.NO_BIAS_DECAY:
                    optimizer.param_groups[1]['initial_lr'] *= scheduler_dict['gamma'] ** step_count

            if last_epoch != -1:
                scheduler.step()
        ###############################################################################################
        # for dirname, _, filenames in os.walk(DATASET_PATH):
        #     for filename in filenames:
        #         print(os.path.join(dirname, filename))

        # if preprocessing possible
        preprocess_type = config.DATA.PREPROCESS
        cv2_size = (config.DATA.IMG_W, config.DATA.IMG_H)
        if not IS_LOCAL:
            preprocess(os.path.join(DATASET_PATH, 'train', 'train_data', 'NOR'), os.path.join(preprocess_type, 'NOR'), preprocess_type, cv2_size)
            preprocess(os.path.join(DATASET_PATH, 'train', 'train_data', 'AMD'), os.path.join(preprocess_type, 'AMD'), preprocess_type, cv2_size)
            preprocess(os.path.join(DATASET_PATH, 'train', 'train_data', 'RVO'), os.path.join(preprocess_type, 'RVO'), preprocess_type, cv2_size)
            preprocess(os.path.join(DATASET_PATH, 'train', 'train_data', 'DMR'), os.path.join(preprocess_type, 'DMR'), preprocess_type, cv2_size)
            data_dir = preprocess_type
            # data_dir = os.path.join(DATASET_PATH, 'train/train_data')
        else:  # IS_LOCAL
            data_dir = os.path.join(DATASET_PATH, preprocess_type)

        # eda
        # train_std(data_dir, preprocess_type, cv2_size)

        fold_df = split_cv(data_dir, n_splits=config.NUM_FOLDS)
        val_fold_idx = config.IDX_FOLD

        ###############################################################################################

        train_loader = get_dataloader(config, data_dir, fold_df, val_fold_idx, 'train', transform=Albu())
        val_loader = get_dataloader(config, data_dir, fold_df, val_fold_idx, 'val')

        postfix = dict()
        num_epochs = config.TRAIN.NUM_EPOCHS

        val_acc_list = []
        for epoch in range(last_epoch+1, num_epochs):

            if epoch >= config.LOSS.FINETUNE_EPOCH:
                criterion = get_loss(config.LOSS.FINETUNE_LOSS)
            else:
                criterion = get_loss(config.LOSS.NAME)

            train_values = train_single_epoch(config, model, train_loader, criterion, optimizer, scheduler, epoch)
            val_values = evaluate_single_epoch(config, model, val_loader, criterion, epoch)
            val_acc_list.append((epoch, val_values[2]))

            if config.SCHEDULER.NAME != 'one_cyle_lr':
                scheduler.step()

            if IS_LOCAL:
                utils.checkpoint.save_checkpoint(config, model, epoch, val_values[1], val_values[0])

            else:
                postfix['train_loss'] = train_values[0]
                postfix['train_res'] = train_values[1]
                postfix['train_acc'] = train_values[2]
                postfix['train_sens'] = train_values[3]
                postfix['train_spec'] = train_values[4]

                postfix['val_loss'] = val_values[0]
                postfix['val_res'] = val_values[1]
                postfix['val_acc'] = val_values[2]
                postfix['val_sens'] = val_values[3]
                postfix['val_spec'] = val_values[4]

                nsml.report(**postfix, summary=True, step=epoch)

                val_res = '%.10f' % val_values[1]
                val_res = val_res.replace(".", "")
                val_res = val_res[:4] + '.' + val_res[4:8] + '.' + val_res[8:]
                save_name = 'epoch_%02d_score%s_loss%.4f.pth' % (epoch, val_res, val_values[0])
                # nsml.save(save_name)
                nsml.save(epoch)

        for e, val_acc in val_acc_list:
            print('%02d %s' % (e, val_acc))
Ejemplo n.º 31
0
  def run(self):
    mb_obs, mb_td_targets, mb_base_actions, \
    mb_xy0, mb_xy1, \
    mb_values, mb_dones \
      = [], [], [], [], [], [], []

    mb_states = self.states
    for n in range(self.nsteps):
      # pi, pi2, x1, y1, x2, y2, v0
      pi1, pi_xy0, pi_xy1, values, states = self.model.step(
          self.obs, self.states, self.dones)

      pi1_noise = np.random.random_sample((self.nenv, 3)) * 0.3

      base_actions = np.argmax(
          pi1 * self.base_act_mask + pi1_noise, axis=1)
      xy0 = np.argmax(pi_xy0, axis=1)

      x0 = (xy0 % 32).astype(int)
      y0 = (xy0 / 32).astype(int)

      xy1 = np.argmax(pi_xy1, axis=1)
      x1 = (xy1 % 32).astype(int)
      y1 = (xy1 / 32).astype(int)

      # Scripted Agent Hacking

      for env_num in range(self.nenv):
        if env_num >= self.nscripts:  # only for scripted agents
          continue

        ob = self.obs[env_num, :, :, :]
        player_relative = ob[:, :, -1]

        self.group_list[env_num] = common.update_group_list2(
            self.control_groups[env_num])

        if len(self.action_queue[env_num]) == 0:

          self.action_queue[env_num], self.group_id[env_num], self.dest_per_marine[env_num], self.xy_per_marine[env_num] = \
            common.solve_tsp(player_relative,
                             self.selected[env_num][0],
                             self.group_list[env_num],
                             self.group_id[env_num],
                             self.dest_per_marine[env_num],
                             self.xy_per_marine[env_num])

        base_actions[env_num] = 0
        x0[env_num] = 0
        y0[env_num] = 0
        x1[env_num] = 0
        y1[env_num] = 0

        if len(self.action_queue[env_num]) > 0:
          action = self.action_queue[env_num].pop(0)
          base_actions[env_num] = action.get("base_action", 0)

          x0[env_num] = action.get("x0", 0)
          y0[env_num] = action.get("y0", 0)
          xy0[env_num] = y0[env_num] * 32 + x0[env_num]

          x1[env_num] = action.get("x1", 0)
          y1[env_num] = action.get("y1", 0)
          xy1[env_num] = y1[env_num] * 32 + x1[env_num]

      base_actions = self.valid_base_action(base_actions)
      new_base_actions = self.trans_base_actions(base_actions)

      base_action_spec = self.env.action_spec(new_base_actions)
      # print("base_actions:", base_actions)
      actions = self.construct_action(
          base_actions,
          base_action_spec,
          x0,
          y0,
          x1,
          y1
      )

      mb_obs.append(np.copy(self.obs))
      mb_base_actions.append(base_actions)

      mb_xy0.append(xy0)
      mb_xy1.append(xy1)
      mb_values.append(values)
      mb_dones.append(self.dones)

      #print("final acitons : ", actions)
      obs, rewards, dones,\
      available_actions, army_counts,\
      control_groups, selected, xy_per_marine\
      = self.env.step(
          actions=actions)
      self.army_counts = army_counts
      self.control_groups = control_groups
      self.selected = selected
      for env_num, data in enumerate(xy_per_marine):
        self.xy_per_marine[env_num] = data
      self.update_available(available_actions)

      self.states = states
      self.dones = dones
      mean_100ep_reward_a2c = 0
      for n, done in enumerate(dones):
        self.total_reward[n] += float(rewards[n])
        if done:
          self.obs[n] = self.obs[n] * 0
          self.episodes += 1
          num_episodes = self.episodes
          self.episode_rewards.append(self.total_reward[n])

          model = self.model
          mean_100ep_reward = round(
              np.mean(self.episode_rewards[-101:]), 1)
          if (n < self.nscripts):  # scripted agents
            self.episode_rewards_script.append(
                self.total_reward[n])
            mean_100ep_reward_script = round(
                np.mean(self.episode_rewards_script[-101:]), 1)
            nsml.report(
                reward_script=self.total_reward[n],
                mean_reward_script=mean_100ep_reward_script,
                reward=self.total_reward[n],
                mean_100ep_reward=mean_100ep_reward,
                episodes=self.episodes,
                step=self.episodes,
                scope=locals()
            )
          else:
            self.episode_rewards_a2c.append(self.total_reward[n])
            mean_100ep_reward_a2c = round(
                np.mean(self.episode_rewards_a2c[-101:]), 1)
            nsml.report(
                reward_a2c=self.total_reward[n],
                mean_reward_a2c=mean_100ep_reward_a2c,
                reward=self.total_reward[n],
                mean_100ep_reward=mean_100ep_reward,
                episodes=self.episodes,
                step=self.episodes,
                scope=locals()
            )
            print("mean_100ep_reward_a2c", mean_100ep_reward_a2c)

          if self.callback is not None:
            self.callback(locals(), globals())
          self.total_reward[n] = 0
          self.group_list[n] = []


      self.update_obs(obs)
      mb_td_targets.append(rewards)
    mb_dones.append(self.dones)
    #batch of steps to batch of rollouts
    mb_obs = np.asarray(
        mb_obs, dtype=np.uint8).swapaxes(1, 0).reshape(
            self.batch_ob_shape)
    mb_td_targets = np.asarray(mb_td_targets, dtype=np.float32).swapaxes(1, 0)
    mb_base_actions = np.asarray(
        mb_base_actions, dtype=np.int32).swapaxes(1, 0)

    mb_xy0 = np.asarray(mb_xy0, dtype=np.int32).swapaxes(1, 0)
    mb_xy1 = np.asarray(mb_xy1, dtype=np.int32).swapaxes(1, 0)

    mb_values = np.asarray(mb_values, dtype=np.float32).swapaxes(1, 0)
    mb_dones = np.asarray(mb_dones, dtype=np.bool).swapaxes(1, 0)
    mb_masks = mb_dones[:, :-1]
    mb_dones = mb_dones[:, 1:]
    last_values = self.model.value(self.obs, self.states,
                                   self.dones).tolist()
    #discount/bootstrap off value fn
    for n, (rewards, dones, value) in enumerate(
        zip(mb_td_targets, mb_dones, last_values)):
      rewards = rewards.tolist()
      dones = dones.tolist()
      if dones[-1] == 0:
        rewards = discount_with_dones(rewards + [value], dones + [0],
                                      self.gamma)[:-1]
      else:
        rewards = discount_with_dones(rewards, dones, self.gamma)
      mb_td_targets[n] = rewards
    mb_td_targets = mb_td_targets.flatten()
    mb_base_actions = mb_base_actions.flatten()
    mb_xy0 = mb_xy0.flatten()
    mb_xy1 = mb_xy1.flatten()

    mb_values = mb_values.flatten()
    mb_masks = mb_masks.flatten()
    return mb_obs, mb_states, mb_td_targets, mb_masks, \
           mb_base_actions, mb_xy0, mb_xy1, mb_values