B #=6\65@sJddlmZddlmZddlmZddlmZddlmZddlmZddl m Z ddl m Z m Z dd lmZdd lZdd lZdd lZdd lZdd lmZmZmZmZd Zd Zd0ddZddZddZd1ddZddZ ddZ!ddZ"ddZ#ddZ$d d!Z%d"d#Z&Gd$d%d%e'Z(d&d'Z)d(d)Z*d2d*d+Z+d,d-Z,d.d/Z-d S)3)absolute_import)division)print_function)unicode_literals) caffe2_pb2)container_abcs) viewitems) DecodeErrorMessage) text_formatN) integer_types binary_type text_type string_typesZoptimizer_iterationZiteration_mutexcs\pgttsgtddDs8tdfdd}||}||}||kS)zM Two ops are identical except for each field in the `ignore_fields`. css|]}t|tVqdS)N) isinstancer).0fr6/tmp/pip-install-l3r2oljg/torch/caffe2/python/utils.py sz OpAlmostEqual..z*Expect each field is text type, but got {}cs0t|}x D]}||r||qW|S)N)copydeepcopyHasFieldZ ClearField)opfield) ignore_fieldsrrclean_op#s    zOpAlmostEqual..clean_op)rlistallAssertionErrorformat)Zop_aZop_brrr)rr OpAlmostEquals   r!cCsP|jdkr0tj|jtjd|j|j|j|jStj|jtjd|j j SdS)Nr)dtype) numnpasarraydatafloat32reshapeZchannelsheightwidthshapeZdim)ZblobrrrCaffeBlobToNumpyArray/s  r,cCs0|jtjjkr(tj|jtjd|j S|jtjj krPtj|j tj d|j S|jtjj krxtj|jtjd|j S|jtjjkrtj|jtjd|j S|jtjjkrtj|jtjd|j S|jtjjkrtj|jtjd|j S|jtjjkrtj|jtjd|j Stdt|jdS)N)r"z$Tensor data type not supported yet: ) data_typer TensorProtoFLOATr$r% float_datar'r(dimsDOUBLE double_datafloat64INT32 int32_dataintINT16int16UINT16uint16INT8int8UINT8uint8 RuntimeErrorstr)tensorrrrCaffe2TensorToNumpyArray:s.rCcCst}|j|j|r ||_|jtjkrTtjj |_ |j t | tn^|jtjkrtjj|_ |jt | tjn(|jtjks|jtjkrtjj|_ |j| tjn|jtjkrtjj|_ |jt | tjn|jtjkr4tjj|_ |jt | tjn~|jtjkrjtjj|_ |jt | tjnH|jtjkrtjj|_ |jt | tjnt dt!|j|S)Nz#Numpy data type not supported yet: )"rr.r1extendr+namer"r$r'r/r-r0rflattenZastypefloatr4r2r3r7Zint32r5r6tolistr9r8r;r:r=r<r?r>r@rA)ZarrrErBrrrNumpyArrayToCaffe2TensorVs8           rIc CsRt}||_t|tj}t|tjrL|jj tj krL|j | |St|tjrf| }nt|tjr|t|}t |tkr||_nt |tkst |tkr||_nt|tr||_nt|tr|d|_nlt|tjr|j|nPt|tr||_n6|rLtdd|DrL|j dd|Dn|r~tdd|Dr~|j dd|Dn|rtdd|Dr|j! dd|Dn|rtd d|Dr|j" |nv|r td d|Dr |j! d d|DnD|r8t#d $||t |t%d d|Dnt#d$||t ||S)z*Makes an argument based on the value type.zutf-8css |]}t|ttjgkVqdS)N)typerGr$float_)rvrrrrszMakeArgument..css(|] }t|tjkr|n|VqdS)N)rJr$rKitem)rrLrrrrscss,|]$}t|tkp"t|ttjgkVqdS)N)rJr boolr$int_)rrLrrrrscss(|] }t|tjkr|n|VqdS)N)rJr$rOrM)rrLrrrrscss"|]}t|tpt|tVqdS)N)rr r)rrLrrrrscss&|]}t|tr|dn|VqdS)zutf-8N)rrencode)rrLrrrrscss|]}t|tjVqdS)N)rrNetDef)rrLrrrrscss|]}t|tVqdS)N)rr )rrLrrrrscss|]}|VqdS)N)SerializeToString)rrLrrrrszBUnknown iterable argument type: key={} value={}, value type={}[{}]css|]}t|VqdS)N)rJ)rrLrrrrsz5Unknown argument type: key={} value={}, value type={})&rZArgumentrErrIterabler$Zndarrayr"rJr'ZfloatsrDrFrHZgenericZasscalarrGrr rNir srrPrQnZCopyFromr rRrintsstringsnets ValueErrorr set)keyvalueargumentiterablerrr MakeArgumentws\             "r`cCs<|}yt|||Stjk r6|||SXdS)a,Reads a protobuffer with the given proto class. Inputs: cls: a protobuffer class. s: a string of either binary or text protobuffer content. Outputs: proto: the protobuffer of cls Throws: google.protobuf.message.DecodeError: if we cannot decode the message. N)r Parse ParseErrorZParseFromString)clsrUobjrrrTryReadProtoWithClasss   recCs.x(t|D]\}}t||kr ||Sq WdS)zNGets a specific field from a protocol buffer that matches the given class N)rrJ)rd function_maprcfuncrrrGetContentFromProtos rhc CsLxFt|D]2\}}yt||}||Stk r:w Yq Xq WtddS)Nz$Cannot find a fit protobuffer class.)rrer )rUrfrcrgrdrrrGetContentFromProtoStrings  ric Cs:t|t|}t|d}||WdQRXdS)z:Convert a text file of the given protobuf class to binary.wN)reopenreadwriterR)Z proto_classfilenameZ out_filenameprotoZfidrrrConvertProtoToBinarys rpc Csddlm}m}||jdgdg||jdd|d}|dddf|dddft |dddft |dddfdS) z|Get GPU memory usage stats from CUDAContext/HIPContext. This requires flag --caffe2_gpu_memory_tracking to be enabledr) workspacecoreZGetGPUMemoryUsageZ ____mem____) device_optionN)Z total_by_gpuZ max_by_gputotalZ max_total) caffe2.pythonrqrrRunOperatorOnceCreateOperator DeviceOptionZ GpuDeviceTypeZ FetchBlobr$sum)rqrrbrrrGetGPUMemoryUsageStatss r|c Cs<ddlm}m}||jdt|t||tjddS)Nr)rqrrZFree)rs) rvrqrrrwrxrryrCPU)Zblobsrqrrrrr ResetBlobssr~c@seZdZdZeddZdS) DebugModea= This class allows to drop you into an interactive debugger if there is an unhandled exception in your python script Example of usage: def main(): # your code here pass if __name__ == '__main__': from caffe2.python.utils import DebugMode DebugMode.run(main) cCsly|Stk rYnLtk rfddl}tdttdt|tdYnXdS)NrzjEntering interactive debugger. Type "bt" to print the full stacktrace. Type "help" to see command listing.rt)KeyboardInterrupt Exceptionpdbprintsysexc_infoZ post_mortemexit)rcrgrrrrrun s z DebugMode.runN)__name__ __module__ __qualname____doc__ classmethodrrrrrrsrcCs||krtd|||dS)Nz {}. {} != {})rr )ar{msgrrrraiseIfNotEqual4srcstfdd}|S)z Use this method to decorate your function with DebugMode's functionality Example: @debug def test_foo(self): raise Exception("Bar") csfdd}t|S)Ncs S)Nrr)argsrkwargsrrrgGsz$debug..wrapper..func)rr)rrrg)r)rrrwrapperEszdebug..wrapper) functoolswraps)rrr)rrdebug9s rc Cs|dk r |nt}|dk r|nt}ddlm}||s||tj@|j g|dg||j j d}| g|g}| ||g|gWdQRXn ||}|S)z Often, a mutex guarded iteration counter is needed. This function creates a mutex iter in the net uniquely (if the iter already existing, it does nothing) This function returns the iter blob Nr)rrrt)r+r]r")OPTIMIZER_ITERATION_NAMEITERATION_MUTEX_NAMErvrrZ BlobIsDefinedZ DeviceScoperyrr}Z ConstantFillZDataTypeZINT64Z CreateMutexZ AtomicIterZ GetBlobRef)Zinit_netnetiterZ iter_mutexZiter_valrr iterationrrrBuildUniqueMutexIterNs    rc st|tkstixnt|D]b}||krt||ttrksxtd||fddDd|<qWS)Nz|sz$EnumClassKeyVals..r) rJrdiruppergetattrrrvaluesr )rckr)rrLrEnumClassKeyValsps   $ rcCsi}x|D]}|dsq x|jjD]n}|jdkr4q$|j|jkrb||jrbt||j||j<Pq$|j|jkr$t||j}t|dkr$|||j<Pq$Wd||j<q W|S)z Convert a list of arguments to a name, value dictionary. Assumes that each argument has a name. Otherwise, the argument is skipped. rErN) rZ DESCRIPTORfieldsrElabelZLABEL_OPTIONALrZLABEL_REPEATEDlen)rZansargdlist_rrr ArgsToDicts"       r)N)N)NNr). __future__rrrrZ caffe2.protorZcaffe2.python.compatibilityrZ future.utilsrZgoogle.protobuf.messager r Zgoogle.protobufr rrrZnumpyr$sixr r rrrrr!r,rCrIr`rerhrirpr|r~objectrrrrrrrrrrsD           !H  $