def train(args):
    workspace = cfg.workspace
    te_fold = cfg.te_fold
    n_events = args.n_events
    snr = args.snr

    feature_dir = os.path.join(workspace, "features", "logmel",
                               "n_events=%d" % n_events)
    yaml_dir = os.path.join(workspace, "mixed_audio", "n_events=%d" % n_events)
    (tr_x, tr_at_y, tr_sed_y, tr_na_list, te_x, te_at_y, te_sed_y,
     te_na_list) = pp_data.load_data(feature_dir=feature_dir,
                                     yaml_dir=yaml_dir,
                                     te_fold=te_fold,
                                     snr=snr,
                                     is_scale=is_scale)

    print(tr_x.shape, tr_at_y.shape)
    print(te_x.shape, te_at_y.shape)
    (_, n_time, n_freq) = tr_x.shape
    n_out = len(cfg.events)

    if False:
        for e in tr_x:
            plt.matshow(e.T, origin='lower', aspect='auto')
            plt.show()

    # Build model.
    lay_in = InputLayer(in_shape=(n_time, n_freq))

    a = Reshape((1, n_time, n_freq))(lay_in)
    a = Conv2D(n_outfmaps=64,
               n_row=3,
               n_col=5,
               act='linear',
               strides=(1, 1),
               border_mode=(1, 2))(a)
    a = BN(axis=(0, 2, 3))(a)
    a = Activation('relu')(a)
    a = Conv2D(n_outfmaps=64,
               n_row=3,
               n_col=5,
               act='linear',
               strides=(1, 1),
               border_mode=(1, 2))(a)
    a = BN(axis=(0, 2, 3))(a)
    a = Activation('relu')(a)
    a = Dropout(p_drop=0.2)(a)

    a = Conv2D(n_outfmaps=64,
               n_row=3,
               n_col=5,
               act='linear',
               strides=(1, 1),
               border_mode=(1, 2))(a)
    a = BN(axis=(0, 2, 3))(a)
    a = Activation('relu')(a)
    a = Conv2D(n_outfmaps=64,
               n_row=3,
               n_col=5,
               act='linear',
               strides=(1, 1),
               border_mode=(1, 2))(a)
    a = BN(axis=(0, 2, 3))(a)
    a = Activation('relu')(a)
    a = Dropout(p_drop=0.2)(a)

    a = Conv2D(n_outfmaps=64,
               n_row=3,
               n_col=5,
               act='linear',
               strides=(1, 1),
               border_mode=(1, 2))(a)
    a = BN(axis=(0, 2, 3))(a)
    a = Activation('relu')(a)
    a = Conv2D(n_outfmaps=64,
               n_row=3,
               n_col=5,
               act='linear',
               strides=(1, 1),
               border_mode=(1, 2))(a)
    a = BN(axis=(0, 2, 3))(a)
    a = Activation('relu')(a)
    a = Dropout(p_drop=0.2)(a)

    a = Conv2D(n_outfmaps=n_out,
               n_row=1,
               n_col=1,
               act='sigmoid',
               border_mode=(0, 0),
               name='seg_masks')(a)

    a8 = Lambda(_global_avg_pooling, name='a8')(a)

    md = Model([lay_in], [a8])
    md.compile()
    md.summary(is_logging=True)

    # Callbacks.
    md_dir = os.path.join(workspace, "models", pp_data.get_filename(__file__),
                          "n_events=%d" % n_events, "fold=%d" % te_fold,
                          "snr=%d" % snr)
    pp_data.create_folder(md_dir)
    save_model = SaveModel(md_dir, call_freq=50, type='iter', is_logging=True)
    validation = Validation(te_x=te_x,
                            te_y=te_at_y,
                            batch_size=50,
                            call_freq=50,
                            metrics=['binary_crossentropy'],
                            dump_path=None,
                            is_logging=True)

    callbacks = [save_model, validation]

    observe_nodes = [md.find_layer('seg_masks').output_]
    f_forward = md.get_observe_forward_func(observe_nodes)

    # Generator.
    tr_gen = DataGenerator(batch_size=32, type='train')
    eva_gen = DataGenerator2(batch_size=32, type='test')

    # Train.
    loss_ary = []
    t1 = time.time()
    optimizer = Adam(1e-3)
    for (batch_x, batch_y) in tr_gen.generate(xs=[tr_x], ys=[tr_at_y]):
        if md.iter_ % 50 == 0:
            logging.info("iter: %d tr_loss: %f time: %s" % (
                md.iter_,
                np.mean(loss_ary),
                time.time() - t1,
            ))
            t1 = time.time()
            loss_ary = []
        # if md.iter_ % 200 == 0:
        # write_out_at_sed(md, eva_gen, f_forward, te_x, te_at_y, te_sed_y, n_events, snr, te_fold)
        if md.iter_ == 5001:
            break
        loss = md.train_on_batch(batch_x,
                                 batch_y,
                                 loss_func='binary_crossentropy',
                                 optimizer=optimizer,
                                 callbacks=callbacks)
        loss_ary.append(loss)
### Build model

lay_in = InputLayer(in_shape=(n_in, ))
a = Dense(n_out=n_hid, act='relu', name='dense1')(lay_in)
a = Dropout(p_drop=0.2)(a)
a = Dense(n_out=n_hid, act='relu', name='dense2')(a)
a = Dropout(p_drop=0.2)(a)
lay_out = Dense(n_out=n_out, act='softmax')(a)

md = Model(in_layers=[lay_in], out_layers=[lay_out])
md.compile()
md.summary()

# observe forward
observe_nodes = [
    md.find_layer('dense1').output_,
    md.find_layer('dense2').output_
]
f_forward = md.get_observe_forward_func(observe_nodes)
print md.run_function(func=f_forward, z=[te_x], batch_size=500, tr_phase=0.)

# observe backward
md.set_gt_nodes(target_dim_list=[2])
loss_node = obj.categorical_crossentropy(md.out_nodes_[0], md.gt_nodes_[0])
gparams = K.grad(loss_node + md.reg_value_, md.params_)
observe_nodes = gparams
f_backward = md.get_observe_backward_func(observe_nodes)
print md.run_function(func=f_backward,
                      z=[te_x, te_y],
                      batch_size=500,
                      tr_phase=0.)
def train():
    _loss_func = _jdc_loss_func0
    
    # load data
    t1 = time.time()
    dict = cPickle.load( open( cfg.scrap_fd+'/denoise_enhance_pool_fft_all0.p', 'rb' ) )
    tr_X, tr_mask, tr_y, tr_na_list, te_X, te_mask, te_y, te_na_list = dict['tr_X'], dict['tr_mask'], dict['tr_y'], dict['tr_na_list'], dict['te_X'], dict['te_mask'], dict['te_y'], dict['te_na_list']
    t2 = time.time()
    
    tr_X = pp_data.wipe_click( tr_X, tr_na_list )
    te_X = pp_data.wipe_click( te_X, te_na_list )
    
    # balance data
    tr_X, tr_mask, tr_y = pp_data.BalanceData2( tr_X, tr_mask, tr_y )
    te_X, te_mask, te_y = pp_data.BalanceData2( te_X, te_mask, te_y )
    
    
    print tr_X.shape, tr_y.shape, te_X.shape, te_y.shape
    [n_songs, n_chunks, n_freq] = te_X.shape
    
    tr_y = tr_y.reshape( (len(tr_y), 1) )
    te_y = te_y.reshape( (len(te_y), 1) )
    
    
    # jdc model
    # classifier
    lay_z0 = InputLayer( (n_chunks,) )          # shape:(n_songs, n_chunks) keep the length of songs
    
    lay_in0 = InputLayer( (n_chunks, n_freq), name='in0' )   # shape: (n_songs, n_chunk, n_freq)
    lay_a1 = lay_in0
    # lay_a1 = Lambda( _conv2d )( lay_a1 )
    
    lay_a1 = Lambda( _reshape_3d_to_4d )( lay_a1 )
    lay_a1 = Convolution2D( 32, 3, 3, act='relu', init_type='glorot_uniform', border_mode=(1,1), strides=(1,1), name='a11' )( lay_a1 )
    lay_a1 = Dropout( 0.2 )( lay_a1 )
    lay_a1 = MaxPool2D( pool_size=(1,2) )( lay_a1 )
    
    lay_a1 = Convolution2D( 64, 3, 3, act='relu', init_type='glorot_uniform', border_mode=(1,1), strides=(1,1), name='a12' )( lay_a1 )
    lay_a1 = Dropout( 0.2 )( lay_a1 )
    lay_a1 = MaxPool2D( pool_size=(1,2) )( lay_a1 )
    lay_a1 = Lambda( _reshape_4d_to_3d )( lay_a1 )
    
    lay_a1 = Dense( n_hid, act='relu', name='a2' )( lay_a1 )       # shape: (n_songs, n_chunk, n_hid)
    lay_a1 = Dropout( 0.2 )( lay_a1 )
    lay_a1 = Dense( n_hid, act='relu', name='a4' )( lay_a1 )
    lay_a1 = Dropout( 0.2 )( lay_a1 )
    lay_a1 = Dense( n_hid, act='relu', name='a6' )( lay_a1 )
    lay_a1 = Dropout( 0.2 )( lay_a1 )
    lay_a8 = Dense( n_out, act='sigmoid', init_type='zeros', b_init=0, name='a8' )( lay_a1 )     # shape: (n_songs, n_chunk, n_out)
    
    # detector
    lay_b1 = lay_in0     # shape: (n_songs, n_chunk, n_freq)
    lay_b2 = Lambda( _conv2d )( lay_b1 )    # shape: (n_songs, n_chunk, n_freq)
    lay_b2 = Lambda( _reshape_3d_to_4d )( lay_b1 )
    lay_b2 = MaxPool2D( pool_size=(1,2) )( lay_b2 )
    lay_b2 = Lambda( _reshape_4d_to_3d )( lay_b2 )
    lay_b8 = Dense( n_out, act='hard_sigmoid', init_type='zeros', b_init=-2.3, name='b8' )( lay_b2 )
    md = Model( in_layers=[lay_in0, lay_z0], out_layers=[lay_a8, lay_b8], any_layers=[] )
    
      
    # print summary info of model
    md.summary()

    # callbacks (optional)
    # save model every n epoch (optional)
    pp_data.CreateFolder( cfg.wbl_dev_md_fd )
    pp_data.CreateFolder( cfg.wbl_dev_md_fd+'/cnn_fft' )
    save_model = SaveModel( dump_fd=cfg.wbl_dev_md_fd+'/cnn_fft', call_freq=20, type='iter' )
    validation = Validation( tr_x=None, tr_y=None, va_x=None, va_y=None, te_x=[te_X, te_mask], te_y=te_y, batch_size=100, metrics=[_loss_func], call_freq=20, dump_path=None, type='iter' )
    
    # callbacks function
    callbacks = [save_model, validation]

    
    # EM training
    md.set_gt_nodes( tr_y )
    md.find_layer('a11').set_trainable_params( ['W','b'] )
    md.find_layer('a12').set_trainable_params( ['W','b'] )
    md.find_layer('a2').set_trainable_params( ['W','b'] )
    md.find_layer('a4').set_trainable_params( ['W','b'] )
    md.find_layer('a6').set_trainable_params( ['W','b'] )
    md.find_layer('a8').set_trainable_params( ['W','b'] )
    md.find_layer('b8').set_trainable_params( [] )
    opt_classifier = Adam( 1e-3 )
    f_classify = md.get_optimization_func( loss_func=_loss_func, optimizer=opt_classifier, clip=None )
    
    
    md.find_layer('a11').set_trainable_params( [] )
    md.find_layer('a12').set_trainable_params( [] )
    md.find_layer('a2').set_trainable_params( [] )
    md.find_layer('a4').set_trainable_params( [] )
    md.find_layer('a6').set_trainable_params( [] )
    md.find_layer('a8').set_trainable_params( [] )
    md.find_layer('b8').set_trainable_params( ['W','b'] )
    opt_detector = Adam( 1e-3 )
    f_detector = md.get_optimization_func( loss_func=_loss_func, optimizer=opt_detector, clip=None )
    
    
    _x, _y = md.preprocess_data( [tr_X, tr_mask], tr_y, shuffle=True )
    
    for i1 in xrange(500):
        print '-----------------------'
        opt_classifier.reset()
        md.do_optimization_func_iter_wise( f_classify, _x, _y, batch_size=100, n_iters=80, callbacks=callbacks, verbose=1 )
        print '-----------------------'
        opt_detector.reset()
        md.do_optimization_func_iter_wise( f_detector, _x, _y, batch_size=100, n_iters=20, callbacks=callbacks, verbose=1 )