ó ùµÈ[c@sŽdZddlZddlZddlmZddlmZddlm Z m Z ddl m Z dd l mZd e fd „ƒYZdS( s¨A `BucketingModule` implement the `BaseModule` API, and allows multiple symbols to be used depending on the `bucket_key` provided by each different mini-batch of data. iÿÿÿÿNi(tcontext(tUniformi(t BaseModulet_check_input_names(tModule(t NameManagertBucketingModulec BsŽeZdZdeejƒdddddd„Zd„Zd„Z e d„ƒZ e d„ƒZ e d„ƒZ e d„ƒZe d„ƒZd „Zeeed „Zed ƒddeeed „Zed „Zddd„Zdeeeddd„Zdd„Zddd!ed„Zdd„Zdd„Zdd„Zd„Zed„Zed„Z ed„Z!e d„ƒZ"d„Z#RS("skThis module helps to deal efficiently with varying-length inputs. Parameters ---------- sym_gen : function A function when called with a bucket key, returns a triple ``(symbol, data_names, label_names)``. default_bucket_key : str (or any python object) The key for the default bucket. logger : Logger context : Context or list of Context Defaults to ``mx.cpu()`` work_load_list : list of number Defaults to ``None``, indicating uniform workload. fixed_param_names: list of str Defaults to ``None``, indicating no network parameters are fixed. state_names : list of str States are similar to data and label, but not provided by data iterator. Instead they are initialized to 0 and can be set by set_states() group2ctxs : dict of str to context or list of context, or list of dict of str to context Default is `None`. Mapping the `ctx_group` attribute to the context assignment. compression_params : dict Specifies type of gradient compression and additional arguments depending on the type of compression being used. For example, 2bit compression requires a threshold. Arguments would then be {'type':'2bit', 'threshold':0.5} See mxnet.KVStore.set_gradient_compression method for more details on gradient compression. c Cs€tt|ƒjd|ƒ|dk s+t‚||_||_|j|ƒ\} } } | dk rmt| ƒng} | dk r‹t| ƒng} |dk r©t|ƒng}|dk rÇt|ƒng}t | | dt ƒt | | dt ƒt | |dt ƒt | |dt ƒ| |_ ||_ ||_||_||_||_i|_d|_d|_t |_d|_dS(Ntloggertdatatlabeltstatet fixed_param(tsuperRt__init__tNonetAssertionErrort_default_bucket_keyt_sym_gent _call_sym_gentlistRtTruetFalset_compression_paramst_fixed_param_namest _state_namest_contextt_work_load_listt _group2ctxst_bucketst _curr_modulet_curr_bucket_keyt _params_dirtyt_monitor( tselftsym_gentdefault_bucket_keyRRtwork_load_listtfixed_param_namest state_namest group2ctxstcompression_paramstsymbolt data_namest label_names((s]/usr/local/lib/python2.7/site-packages/mxnet-1.3.1-py2.7.egg/mxnet/module/bucketing_module.pyR As0            cCs(t|_i|_d|_d|_dS(s+Internal utility function to reset binding.N(RtbindedRRRR(R!((s]/usr/local/lib/python2.7/site-packages/mxnet-1.3.1-py2.7.egg/mxnet/module/bucketing_module.pyt _reset_bindcs   cOs$tƒ|j||ŽSWdQXdS(N(RR(R!targstkwargs((s]/usr/local/lib/python2.7/site-packages/mxnet-1.3.1-py2.7.egg/mxnet/module/bucketing_module.pyRjs cCs6|jr|jjS|j|jƒ\}}}|SdS(s1A list of names for data required by this module.N(R,RR*RR(R!t_R*((s]/usr/local/lib/python2.7/site-packages/mxnet-1.3.1-py2.7.egg/mxnet/module/bucketing_module.pyR*ns  cCs<|jr|jjS|j|jƒ\}}}|jƒSdS(s/A list of names for the outputs of this module.N(R,Rt output_namesRRt list_outputs(R!R)R0((s]/usr/local/lib/python2.7/site-packages/mxnet-1.3.1-py2.7.egg/mxnet/module/bucketing_module.pyR1ws  cCs|jst‚|jjS(scGet data shapes. Returns ------- A list of `(name, shape)` pairs. (R,RRt data_shapes(R!((s]/usr/local/lib/python2.7/site-packages/mxnet-1.3.1-py2.7.egg/mxnet/module/bucketing_module.pyR3€scCs|jst‚|jjS(s.Get label shapes. Returns ------- A list of `(name, shape)` pairs. The return value could be ``None`` if the module does not need labels, or if the module is not bound for training (in this case, label information is not available). (R,RRt label_shapes(R!((s]/usr/local/lib/python2.7/site-packages/mxnet-1.3.1-py2.7.egg/mxnet/module/bucketing_module.pyR4‹s cCs|jst‚|jjS(sfGets output shapes. Returns ------- A list of `(name, shape)` pairs. (R,RRt output_shapes(R!((s]/usr/local/lib/python2.7/site-packages/mxnet-1.3.1-py2.7.egg/mxnet/module/bucketing_module.pyR5™scCsC|jr|jst‚|j|j_|jjƒ}t|_|S(s¸Gets current parameters. Returns ------- `(arg_params, aux_params)` A pair of dictionaries each mapping parameter names to NDArray values. (R,tparams_initializedRRRt get_paramsR(R!tparams((s]/usr/local/lib/python2.7/site-packages/mxnet-1.3.1-py2.7.egg/mxnet/module/bucketing_module.pyR7¤s  c Cs”|s2|jddd|d|d|d|ƒdS|jrY| rYtjddd ƒdS|jj||d|d|d |ƒt|_t|_dS( sAssigns parameters and aux state values. Parameters ---------- arg_params : dict Dictionary of name to value (`NDArray`) mapping. aux_params : dict Dictionary of name to value (`NDArray`) mapping. allow_missing : bool If true, params could contain missing values, and the initializer will be called to fill those missing params. force_init : bool If true, will force re-initialize even if already initialized. allow_extra : boolean, optional Whether allow extra parameters that are not needed by symbol. If this is True, no error will be thrown when arg_params or aux_params contain extra parameters that is not needed by the executor. Examples -------- >>> # An example of setting module parameters. >>> sym, arg_params, aux_params = mx.model.load_checkpoint(model_prefix, n_epoch_load) >>> mod.set_params(arg_params=arg_params, aux_params=aux_params) t initializert arg_paramst aux_paramst allow_missingt force_initNsMParameters already initialized and force_init=False. set_params call ignored.t stacklevelit allow_extra( t init_paramsRR6twarningstwarnRt set_paramsRR(R!R:R;R<R=R?((s]/usr/local/lib/python2.7/site-packages/mxnet-1.3.1-py2.7.egg/mxnet/module/bucketing_module.pyRC²s    g{®Gáz„?c Csp|jr| rdS|js)tdƒ‚|jjd|d|d|d|d|d|ƒt|_t|_dS( sŒInitializes parameters. Parameters ---------- initializer : Initializer arg_params : dict Defaults to ``None``. Existing parameters. This has higher priority than `initializer`. aux_params : dict Defaults to ``None``. Existing auxiliary states. This has higher priority than `initializer`. allow_missing : bool Allow missing values in `arg_params` and `aux_params` (if not ``None``). In this case, missing values will be filled with `initializer`. force_init : bool Defaults to ``False``. allow_extra : boolean, optional Whether allow extra parameters that are not needed by symbol. If this is True, no error will be thrown when arg_params or aux_params contain extra parameters that is not needed by the executor. Ns,call bind before initializing the parametersR9R:R;R<R=R?(R6R,RRR@RRR(R!R9R:R;R<R=R?((s]/usr/local/lib/python2.7/site-packages/mxnet-1.3.1-py2.7.egg/mxnet/module/bucketing_module.pyR@Ýs   cCs+|jr|jst‚|jjd|ƒS(s´Gets states from all devices. Parameters ---------- merge_multi_context : bool Default is `True`. In the case when data-parallelism is used, the states will be collected from multiple devices. A `True` value indicate that we should merge the collected results so that they look like from a single executor. Returns ------- list of NDArrays or list of list of NDArrays If `merge_multi_context` is ``True``, it is like ``[out1, out2]``. Otherwise, it is like ``[[out1_dev1, out1_dev2], [out2_dev1, out2_dev2]]``. All the output elements are `NDArray`. tmerge_multi_context(R,R6RRt get_states(R!RD((s]/usr/local/lib/python2.7/site-packages/mxnet-1.3.1-py2.7.egg/mxnet/module/bucketing_module.pyREýscCs/|jr|jst‚|jj||ƒdS(sgSets value for states. Only one of states & values can be specified. Parameters ---------- states : list of list of NDArrays Source states arrays formatted like ``[[state1_dev1, state1_dev2], [state2_dev1, state2_dev2]]``. value : number A single scalar value for all state arrays. N(R,R6RRt set_states(R!tstatestvalue((s]/usr/local/lib/python2.7/site-packages/mxnet-1.3.1-py2.7.egg/mxnet/module/bucketing_module.pyRFs twritecCsZ|jr|jƒ\}} n|r1|jƒn|jrN|jjdƒdS|dksftdƒ‚||_||_ t |_|j |j ƒ\} } } t | | | d|jd|jd|jd|jd|jd |jd |jƒ} | j||||d td dd |ƒ| |_|j |_| |j|j <|jrV|j|| ƒndS(sŠBinding for a `BucketingModule` means setting up the buckets and binding the executor for the default bucket key. Executors corresponding to other keys are bound afterwards with `switch_bucket`. Parameters ---------- data_shapes : list of (str, tuple) This should correspond to the symbol for the default bucket. label_shapes : list of (str, tuple) This should correspond to the symbol for the default bucket. for_training : bool Default is ``True``. inputs_need_grad : bool Default is ``False``. force_rebind : bool Default is ``False``. shared_module : BucketingModule Default is ``None``. This value is currently not used. grad_req : str, list of str, dict of str to str Requirement for gradient accumulation. Can be 'write', 'add', or 'null' (default to 'write'). Can be specified globally (str) or for each argument (list, dict). bucket_key : str (or any python object) bucket key for binding. by default use the default_bucket_key sAlready bound, ignoring bind()Ns2shared_module for BucketingModule is not supportedRRR$R%R&R'R(t force_rebindt shared_moduletgrad_req(R6R7R-R,RtwarningRRt for_trainingtinputs_need_gradRRRRRRRRRRtbindRRRRRC(R!R3R4RNRORJRKRLR:R;R)R*R+tmodule((s]/usr/local/lib/python2.7/site-packages/mxnet-1.3.1-py2.7.egg/mxnet/module/bucketing_module.pyRP s2             cCs|jstdƒ‚||jkr÷|j|ƒ\}}}t|||d|jd|jd|jd|jd|j d|j d|j ƒ}|j |||j j|j jd td |j|jƒ|jd k rç|j|jƒn||j| NDArray. The resulting dict is used for pulling row_sparse parameters from the kvstore, where the str key is the name of the param, and the value is the row id of the param to pull. tsparse_row_id_fnN( R,R6RRSRt provide_datat provide_labelRTRtprepareR(R!t data_batchR`RStoriginal_bucket_keyR3R4((s]/usr/local/lib/python2.7/site-packages/mxnet-1.3.1-py2.7.egg/mxnet/module/bucketing_module.pyRcŸs    cCsN|jr|jst‚|j|j|j|jƒ|jj|d|ƒdS(sØForward computation. Parameters ---------- data_batch : DataBatch is_train : bool Defaults to ``None``, in which case `is_train` is take as ``self.for_training``. tis_trainN( R,R6RRTRSRaRbRtforward(R!RdRf((s]/usr/local/lib/python2.7/site-packages/mxnet-1.3.1-py2.7.egg/mxnet/module/bucketing_module.pyRg¼s  cCs/|jr|jst‚|jjd|ƒdS(sBackward computation.t out_gradsN(R,R6RRtbackward(R!Rh((s]/usr/local/lib/python2.7/site-packages/mxnet-1.3.1-py2.7.egg/mxnet/module/bucketing_module.pyRiÊscCs;|jr|jr|js!t‚t|_|jjƒdS(s:Updates parameters according to installed optimizer and the gradient computed in the previous forward-backward cycle. When KVStore is used to update parameters for multi-device or multi-machine training, a copy of the parameters are stored in KVStore. Note that for `row_sparse` parameters, this function does update the copy of parameters in KVStore, but doesn't broadcast the updated parameters to all devices / machines. Please call `prepare` to broadcast `row_sparse` parameters with the next batch of data. N(R,R6RXRRRRtupdate(R!((s]/usr/local/lib/python2.7/site-packages/mxnet-1.3.1-py2.7.egg/mxnet/module/bucketing_module.pyRjÏs ! cCs+|jr|jst‚|jjd|ƒS(sÙGets outputs from a previous forward computation. Parameters ---------- merge_multi_context : bool Defaults to ``True``. In the case when data-parallelism is used, the outputs will be collected from multiple devices. A ``True`` value indicate that we should merge the collected results so that they look like from a single executor. Returns ------- list of numpy arrays or list of list of numpy arrays If `merge_multi_context` is ``True``, it is like ``[out1, out2]``. Otherwise, it is like ``[[out1_dev1, out1_dev2], [out2_dev1, out2_dev2]]``. All the output elements are numpy arrays. RD(R,R6RRt get_outputs(R!RD((s]/usr/local/lib/python2.7/site-packages/mxnet-1.3.1-py2.7.egg/mxnet/module/bucketing_module.pyRkÞscCs4|jr|jr|js!t‚|jjd|ƒS(sßGets the gradients with respect to the inputs of the module. Parameters ---------- merge_multi_context : bool Defaults to ``True``. In the case when data-parallelism is used, the outputs will be collected from multiple devices. A ``True`` value indicate that we should merge the collected results so that they look like from a single executor. Returns ------- list of NDArrays or list of list of NDArrays If `merge_multi_context` is ``True``, it is like ``[grad1, grad2]``. Otherwise, it is like ``[[grad1_dev1, grad1_dev2], [grad2_dev1, grad2_dev2]]``. All the output elements are `NDArray`. RD(R,R6RORRtget_input_grads(R!RD((s]/usr/local/lib/python2.7/site-packages/mxnet-1.3.1-py2.7.egg/mxnet/module/bucketing_module.pyRlós!cCs2|jr|jst‚|jj|||ƒdS(sõEvaluates and accumulates evaluation metric on outputs of the last forward computation. Parameters ---------- eval_metric : EvalMetric labels : list of NDArray Typically ``data_batch.label``. N(R,R6RRt update_metric(R!t eval_metrictlabelst pre_sliced((s]/usr/local/lib/python2.7/site-packages/mxnet-1.3.1-py2.7.egg/mxnet/module/bucketing_module.pyRms cCs|jst‚|jjS(s,The symbol of the current bucket being used.(R,RRR)(R!((s]/usr/local/lib/python2.7/site-packages/mxnet-1.3.1-py2.7.egg/mxnet/module/bucketing_module.pyR)scCsC|jst‚||_x$|jjƒD]}|j|ƒq(WdS(s"Installs monitor on all executors N(R,RR RRZRR(R!tmonR_((s]/usr/local/lib/python2.7/site-packages/mxnet-1.3.1-py2.7.egg/mxnet/module/bucketing_module.pyRRs N(s learning_rateg{®Gáz„?((s learning_rateg{®Gáz„?($t__name__t __module__t__doc__RtloggingtctxtcpuR R-RtpropertyR*R1R3R4R5R7RRRCRR@RERFRPRTRYRcRgRiRjRkRlRmR)RR(((s]/usr/local/lib/python2.7/site-packages/mxnet-1.3.1-py2.7.egg/mxnet/module/bucketing_module.pyR$sF          *   >         (RtRuRAtRRvR9Rt base_moduleRRRQRtnameRR(((s]/usr/local/lib/python2.7/site-packages/mxnet-1.3.1-py2.7.egg/mxnet/module/bucketing_module.pyts