o ig @sUddlZddlZddlZddlZddlmZddlmZddlm Z m Z m Z m Z m Z mZmZmZmZmZmZmZmZmZddlZddlmZmZmZmZmZddlmZddl m!Z!dd l"m#Z#dd l$m%Z%m&Z&m'Z'm(Z(m)Z)ed Z*ed Z+ed ddZ,edidZ-ee.e/d<d*ddZ0Gddde e*e+fZ1ej2de de de3de1de4f ddZ5ej2de de de3de1de f ddZ6ede1dZ7de7de e3e e3e ffde7fd d!Z8d"gfde7d#e9de7fd$d%Z:de1d&e e1ge7fde7fd'd(Z;gd)ZZ/de deefd?d@Z0de d5eeddfdAdBZ1de defdCdDZ2de defdEdFZ3de d5eddfdGdHZ4de deefdIdJZ5de d5eddfdKdLZ6de deefdMdNZ7de ddfdOdPZ8de dedfdQdRZ9de d5edddfdSdTZ:dUe;dVede!edUe;de!er?r@rArBc Cs||_|dur tt|}t|d|t|d|| dur| nt|_t|_t||_ t||_ t| |_ t ||_ t ||_tjtjd7_tj|_Wdn1sXwYi|_|D]\}} d|j|<| durw||| qddS)zInitialize a new model.Nr4r5r)r1rr.setattrrr2rr6dictr7r:r<listr8r9r#r0r/r3r;items set_param) selfr1rCr5r=r>r?r@rArBr2valuer*r*r-__init__Us.           zModel.__init__r(cC|jS)zmA list of child layers of the model. You can append to it to add layers but not reassign it. )r8rIr*r*r-r?|z Model.layerscCrLr))r9rMr*r*r-r@sz Model.shimscCrL)zfA dict of the model's attrs. You can write to it to update attrs but not reassign it. )r:rMr*r*r-rArNz Model.attrs.cCt|jS)z8Get the names of registered parameter (including unset).)tupler;keysrMr*r*r- param_nameszModel.param_namescstfddjDS)zHGet the names of parameters with registered gradients (including unset).csg|] }|r|qSr*)has_grad).0r1rMr*r- sz$Model.grad_names..)rPrRrMr*rMr- grad_namesszModel.grad_namescCrO)z9Get the names of registered dimensions (including unset).)rPr7rQrMr*r*r- dim_namesrSzModel.dim_namescCrO)z>Get the names of registered node references (including unset).)rPr<rQrMr*r*r- ref_namesrSzModel.ref_names operatorsccs(|jt|}dV|j|dS)aBind arbitrary binary functions to Python operators, for use in any `Model` instance. Can (and should) be used as a contextmanager. EXAMPLE: with Model.define_operators({">>": chain}): model = Relu(512) >> Relu(512) >> Softmax() N)_context_operatorssetrEreset)clsrZtokenr*r*r-define_operatorss zModel.define_operatorscC$||jvrdS|j|durdSdS)zCheck whether the model has a dimension of a given name. If the dimension is registered but the value is unset, returns None. FNT)r7rIr1r*r*r-has_dim z Model.has_dimcCT||jvrtd|d|jd|j|}|dur(d|d|jd}t||S)z4Retrieve the value of a dimension of the given name.zCannot get dimension ' ' for model ''Nz': value unset)r7KeyErrorr1 ValueErrorrIr1rJerrr*r*r-get_dim  z Model.get_dimF)forcerJrncCs||jvrtd|d|jd|j|}tdd|jD}|duo0||ko0| p0|o0|}|rFd|d|jd|d |}t|||j|<dS) zSet a value for a dimension.zCannot set unknown dimension 'rf'.css|] \}}t|VqdSr))bool)rUxyr*r*r- sz Model.set_dim..NzAttempt to change dimension 'z' from z to )r7rhr1anyr;rGri)rIr1rJrn old_value has_paramsinvalid_changerkr*r*r-set_dims   z Model.set_dimcC||r ||SdS)z=Retrieve the value of a dimension of the given name, or None.N)rcrlrbr*r*r- maybe_get_dimzModel.maybe_get_dimcCra)zCheck whether the model has a weights parameter of the given name. Returns None if the parameter is registered but currently unset. FNT)r;rbr*r*r- has_params zModel.has_paramcCsZ||jvrtd|d|jd|j|j|s%td|d|jd|j|j|S)z%Retrieve a weights parameter by name.zUnknown param: 'rfroz Parameter 'z' has not been allocated yet.)r;rhr1r6r|r3 get_paramrbr*r*r-r}s zModel.get_paramcCry)z.Retrieve a weights parameter by name, or None.N)r|r}rbr*r*r-maybe_get_paramr{zModel.maybe_get_paramcCs6|dur d|j|<dS|j|j||d|j|<dS)z Set a weights parameter's value.NT)r;r6rHr3rIr1rJr*r*r-rHszModel.set_paramcC|j|j|S)z@Check whether the model has a non-zero gradient for a parameter.)r6rTr3rbr*r*r-rTzModel.has_gradcCr)zGet a gradient from the model.)r6get_gradr3rbr*r*r-rrzModel.get_gradcC|j|j||dS)z#Set a gradient value for the model.N)r6set_gradr3rr*r*r-rzModel.set_gradcCry)z%Retrieve a gradient by name, or None.N)rTrrbr*r*r-maybe_get_gradr{zModel.maybe_get_gradcCr)z1Increment the gradient of a parameter by a value.N)r6inc_gradr3rr*r*r-r rzModel.inc_gradcCra)zCheck whether the model has a reference of a given name. If the reference is registered but the value is unset, returns None. FNT)r<rbr*r*r-has_refrdz Model.has_refcCre)z4Retrieve the value of a reference of the given name.zCannot get reference 'rfroNz': value unset.)r<rhr1rirjr*r*r-get_refrmz Model.get_refcCry)z8Retrieve the value of a reference if it exists, or None.N)rrrbr*r*r- maybe_get_ref&r{zModel.maybe_get_refcCs8|dur ||j|<dS||vr||j|<dStd)zSet a value for a reference.Nz)Cannot add reference to node not in tree.)r<walkrirr*r*r-set_ref*s  z Model.set_refXis_traincCs|j|||dS)z~Call the model's `forward` function, returning the output and a callback to compute the gradients via backpropagation.rr4)rIrrr*r*r-__call__3szModel.__call__YcCs8tr t|j|j|||jdur|j|||d|S)zFinish initialization of the model, optionally providing a batch of example input and output data to perform shape inference.N)rr)rgetrr1r4r5)rIrrr*r*r- initialize8s  zModel.initializecCs|j||ddS)apRun the model over a batch of data, returning the output and a callback to complete the backward pass. A tuple (Y, finish_update), where Y is a batch of output data, and finish_update is a callback that takes the gradient with respect to the output and an optimizer function, and returns the gradient with respect to the input. TrrrIrr*r*r- begin_updateAszModel.begin_updatecCs|j||dddS)zCall the model's `forward` function with `is_train=False`, and return only the output, instead of the `(output, callback)` tuple. Frrrrr*r*r-predictJsz Model.predict optimizercCsz|D] }|jD]}||q q|D]$}|jD]}||r9||j|f||||\}}|||qqdS)zUpdate parameters with current gradients. The optimizer is called with each parameter and gradient of the model. N) rr@ finish_updaterRrTr3r}rrH)rIrnodeshimr1paramgradr*r*r-rPs       zModel.finish_updatec csi}|jD]}|j|f}||vr ||||<||||qt'}|jD] }|||q)|j D] }|||q7dVWdn1sOwY|re| D] \}}|||qZdSdS)zContext manager to temporarily set the model's parameters to specified values. The params are a dictionary keyed by model IDs, whose values are arrays of weight values. N) rRr3r}rH contextlib ExitStackr? enter_context use_paramsr@rG) rIr>backupr1keystacklayerrrr*r*r-r_s(     zModel.use_paramsbfsorderrcCs@|dkr|S|dkr|jddS|dkr|jddStd)zIterate out layers of the model. Nodes are returned in breadth-first order by default. Other possible orders are "dfs_pre" (depth-first search in preorder) and "dfs_post" (depth-first search in postorder).rdfs_preF) post_orderdfs_postTz5Invalid order, must be one of: bfs, dfs_pre, dfs_post) _walk_bfs _walk_dfsri)rIrr*r*r-rvs  z Model.walkccsJ|g}t}|D]}t||vrq |t||V||jq dS)z/Iterate out layers of the model, breadth-first.N)r\r3addextendr?)rIqueueseenrr*r*r-rs zModel._walk_bfsrccst}|g}t|j|t|<|s|V|rUz%t|t|d}t||vr;|s-|V||t|j|t|<WntyP|rJ|dV|Ynw|sdSdS)z-Iterate out layers of the model, depth-first.N)rEiterr?r3nextappend StopIterationpop)rIrrr next_childr*r*r-rs*      zModel._walk_dfsrcCs~t|D]}||jvr|j|||jvs qt|}|D]}|jD]}||}|dur;||vr;||dq&q!dS)a]Remove a node from all layers lists, and then update references. References that no longer point to a node within the tree will be set to `None`. For instance, let's say a node has its grandchild as a reference. If the child is removed, the grandchild reference will be left dangling, so will be set to None. N)rFrr?remover\rYrr)rIrchildtreer1refr*r*r- remove_nodes       zModel.remove_node)r5cCst|d|t|d|dS)Nr4r5)rD)rIrCr5r*r*r-replace_callbackss zModel.replace_callbacksoldnewcsjd}t|jddD](}|urd}q fdd|jD|_|jD]}||ur1||q"q |S)zzReplace a node anywhere it occurs within the model. Returns a boolean indicating whether the replacement was made.FrrTcsg|] }|ur n|qSr*r*)rUrrrr*r-rVsz&Model.replace_node..)rFrr8rYrr)rIrrrrr1r*rr- replace_nodes   zModel.replace_nodecCsHi}|D]}|jD]}||}||}||f||j|f<q q|S)zGet non-zero gradients of the model's parameters, as a dictionary keyed by the parameter ID. The values are (weights, gradients) tuples. )rrWr}rr3)rI gradientsrr1rrr*r*r- get_gradientss    zModel.get_gradientsrIcCs|S)z Create a copy of the model, its attributes, and its parameters. Any child layers will also be deep-copied. The copy will receive a distinct `model.id` value. )_copyrMr*r*r-copysz Model.copyrc CsB|duri}i}|jD]}||r||nd||<q g}|jD]%}t||vr6|tt|t|q!||}||t|<||q!g}|j D]$}t||vra|tt |t|qL| } | |t|<|| qLt|j |j |jt |t |jt |j||d} |jD] }| ||| qtt| S)N)r5r>r=rAr?r@)rRr|r}r?r3rrr#rr@rrr1r4r5deepcopyr7r:rWrrr") rIrr>r1 copied_layersr copied_layer copied_shimsr copied_shimcopiedr*r*r-rs@              z Model._copygpu_idcCsHddl}|jj||tWddS1swYdS)z)Transfer the model to a given GPU device.rN)cupy.cuda.devicecudadeviceDevice_to_opsr)rIrcupyr*r*r-to_gpus"z Model.to_gpucCs|tdS)zTransfer the model to CPU.N)rrrMr*r*r-to_cpurz Model.to_cpuc Cs|D];}||_|jD]$}||r|||||||r0|||| |q |j D] }| |j |j q4qdS)z Common method for to_cpu/to_gpu.N)rr2rRr|rH asarray_fr}rTrrr@ to_device device_type device_id)rIr2rr1rr*r*r-rs     z Model._to_opscCs.|}t|jjdd}tt||}t|S)aaSerialize the model to a bytes representation. Models are usually serialized using msgpack, so you should be able to call msgpack.loads() on the data and get back a dictionary with the contents. Serialization should round-trip identically, i.e. the same bytes should result from loading and serializing a model. <) byte_order)to_dictrr2to_numpyrrsrsly msgpack_dumps)rImsg to_numpy_ler*r*r-to_bytes&s  zModel.to_bytespathcCsTt|tr t|n|}|d}||WddS1s#wYdS)zSerialize the model to disk. Most models will serialize to a single file, which should just be the bytes contents of model.to_bytes(). wbN) isinstancestrropenwriter)rIrfile_r*r*r-to_disk3s "z Model.to_diskc Csggggd}t|}ddt|D}t|D][\}}i}g}|jD]#}||s1d||<q%||} | j|vrC|| j||<q%||q%|rRtd|i} |j D]} | | rc| | nd| | <qW|d||j | |dq|D])}i} |j D]\}} z t| | ||| |<WqtyYqw|d| qx|D]}|d d d |jDq|D]'}i}|jD]}||rttt||||<qd||<q|d |q|S) zSerialize the model to a dict representation. Serialization should round-trip identically, i.e. the same dict should result from loading and serializing a model. )nodesrAr>r@cSsi|]\}}|j|qSr*)r3)rUirr*r*r- Ksz!Model.to_dict..NzCannot get references: r)indexr1r=rBrAr@cSsg|]}|qSr*)r)rUrr*r*r-rVisz!Model.to_dict..r>)rFr enumeraterYrrr3rrirXrcrlr1rArGserialize_attr TypeErrorr@rRr|rr rr})rIrr node_to_irrrB invalid_refsr1rr=dimrArJr>r*r*r-r;sP            z Model.to_dict bytes_datacCs$t|}tt|jj|}||S)aeDeserialize the model from a bytes representation. Models are usually serialized using msgpack, so you should be able to call msgpack.loads() on the data and get back a dictionary with the contents. Serialization should round-trip identically, i.e. the same bytes should result from loading and serializing a model. )r msgpack_loadsrrr2asarray from_dict)rIrrr*r*r- from_bytests  zModel.from_bytescCsRt|tr t|n|}|d }|}Wdn1swY||S)zDeserialize the model from disk. Most models will serialize to a single file, which should just be the bytes contents of model.to_bytes(). rbN)rrrrreadr)rIrrrr*r*r- from_disks    zModel.from_diskrcCsxd|vr d}t|t|}t|dt|kr tdt|D]\}}|d|}|d|_|dD]\}}|durG|||q9|dD]\} } | dur]| | dqN| | || qN|d|D]\} }|j | } t | || |} | |j | <qn|d|D]\}}|dur|j |}|||qt|d |D] \}}|j||qq$|S) NrzMTrying to read a Model that was created with an incompatible version of Thincz.Cannot deserialize model: mismatched structurer1r=rBrAr>r@)rQrirFrlenrr1rGrxrrArdeserialize_attrr2rrrHr@r)rIrrkrrrinforrJr ref_indexattr default_value loaded_value param_name shim_bytesr*r*r-rs:       zModel.from_dictTstrictrcCsjt|tr t|n|}|s|sdS|d }|}Wdn1s)wY|j||dS)zCheck whether serialized data on disk is compatible with the model. If 'strict', the function returns False if the model has an attribute already loaded that would be changed. FrNr)rrris_direxistsrrcan_from_bytes)rIrrrrr*r*r- can_from_disks  zModel.can_from_diskcCs2zt|}Wn tyYdSw|j||dS)zCheck whether the bytes data is compatible with the model. If 'strict', the function returns False if the model has an attribute already loaded that would be changed. Fr)rrri can_from_dict)rIrrrr*r*r-r s  zModel.can_from_bytesc Csd|vrdSt|}t|dt|krdSt|D]\}}|d|}|r4|d|jkr4dSt|d|t|jkrDdS|dD]\}}||} | dur[dS| rh| ||krhdSqJ|d|D]&\} }| | } | durdS| r|dur| | } | j |j krdSqq|r|d|D]*\} }| |j vrzt|j | |j | | |}Wn tyYqw||krdSqqd S) zCheck whether a dictionary is compatible with the model. If 'strict', the function returns False if the model has an attribute already loaded that would be changed. rFr1r@r=r>NrAT)rQrFrrrr1r@rGrcrlr|r}shaperArr)rIrrrrrrrrJrcrr|rr serializedr*r*r-r sR          zModel.can_from_dictothercC*d|jvr td|jd||S)z-Apply the function bound to the '+' operator.+zUndefined operator: +r[rrrIrr*r*r-__add__z Model.__add__cCr)z-Apply the function bound to the '-' operator.-zUndefined operator: -rrr*r*r-__sub__rz Model.__sub__cCr)z-Apply the function bound to the '*' operator.*zUndefined operator: *rrr*r*r-__mul__rz Model.__mul__cCr)z-Apply the function bound to the '@' operator.@zUndefined operator: @rrr*r*r- __matmul__rzModel.__matmul__cCrz-Apply the function bound to the '/' operator./zUndefined operator: /rrr*r*r-__div__rz Model.__div__cCrrrrr*r*r- __truediv__ rzModel.__truediv__cCr)z.Apply the function bound to the '//' operator.z//zUndefined operator: //rrr*r*r- __floordiv__rzModel.__floordiv__cCr)z-Apply the function bound to the '%' operator.%zUndefined operator: %rrr*r*r-__mod__rz Model.__mod__cKr)z.Apply the function bound to the '**' operator.z**zUndefined operator: **r)rIrr,r*r*r-__pow__rz Model.__pow__cCr)z.Apply the function bound to the '<<' operator.z<>' operator.z>>zUndefined operator: >>rrr*r*r- __rshift__)rzModel.__rshift__cCr)z-Apply the function bound to the '&' operator.&zUndefined operator: &rrr*r*r-__and__/rz Model.__and__cCr)z-Apply the function bound to the '^' operator.^zUndefined operator: ^rrr*r*r-__xor__5rz Model.__xor__cCr)z-Apply the function bound to the '|' operator.|zUndefined operator: |rrr*r*r-__or__;rz Model.__or__)NN)F)rr#r(Nr))r(N)k__name__ __module__ __qualname____doc__r/int__annotations__ threadingLockr0r%r[rrrrrr r rrrp __slots__rr rrrrKpropertyr?r@rArrRrWrXrY classmethodrcontextmanagerr`rcrlrxrzr|r}r~rHrTrrrrrrrrr r!rrrrrrrrrrrrrrrr"rrrrrbytesrrrrrrrr r r rrrrrrr r"r#r$r%r'r)r+r*r*r*r-r#1s        '       $     &   )   9  $  /_rJr1cC t|S)zSerialize an attribute value (defaults to msgpack). You can register custom serializers using the @serialize_attr.register decorator with the type to serialize, e.g.: @serialize_attr.register(MyCustomObject). )rrr9rJr1r'r*r*r-rB rcCr:)zDeserialize an attribute value (defaults to msgpack). You can register custom deserializers using the @deserialize_attr.register decorator with the type to deserialize, e.g.: @deserialize_attr.register(MyCustomObject). )rrr;r*r*r-rKr<r_ModelTmappingcCsL|D]}|j|vr#||j}|D]\}}||jvr"||j|<qq|S)zWalk over the model's nodes, changing the value of attributes using the provided mapping, which maps node names to attr names to attr values. )rr1rGrA)r'r>rrArrJr*r*r-change_attr_valuesWs     r? dropout_ratedropcCs0|D]}|D] }||jvr||j|<qq|S)zWalk over the model's nodes, setting the dropout rate. You can specify one or more attribute names, by default it looks for ["dropout_rate"]. )rrA)r'rArArrr*r*r-set_dropout_rateds   rBwrappercCs*t|D] }||||q||S)zORecursively wrap a model and its submodules. The model is updated in-place.)rFrr)r'rCrr*r*r-wrap_model_recursiveosrD)r#rrr?rBrD)r'r#r(r#)=rr functoolsr2 contextvarsrpathlibrtypingrrrrrr r r r r rrrrrbackendsrrrrr optimizersrr@rtypesrutilrrrrrr r!r"r%rEr1r.r#singledispatchrr8rrr=r?floatrBrD__all__r*r*r*r-sD   @        &