/usr/local/lib/python3.6/site-packages/transformers/benchmark/__pycache__
Edit: /usr/local/lib/python3.6/site-packages/transformers/benchmark/__pycache__/benchmark.cpython-36.pyc (7545B)
3
Eg) @ s d Z ddlZddlmZmZ ddlmZ ddlmZm Z ddl
m Z m Z m
Z
dd lmZmZmZmZmZmZ e rddlZdd
lmZ e rddljZe
je Z G d d d eZ dS )
zD
Benchmarking the library on inference and training in PyTorch.
N)CallableOptional )PretrainedConfig)
MODEL_MAPPINGMODEL_WITH_LM_HEAD_MAPPING)is_py3nvml_availableis_torch_availablelogging ) BenchmarkMemory
MemorySummarymeasure_peak_memory_cpustart_memory_tracingstop_memory_tracing)PyTorchBenchmarkArgumentsc @ s e Zd ZU eedZee dd Z
ee e e dddZ
ee e eee gdddZee e e dd d
Zee e eee gdd d Zee e eg d
f dddZee e eg d
f dddZe dddZeg d
f eegdddZd
S )PyTorchBenchmarkZPyTorchc C s t jS )N)torch __version__)self r J/usr/local/lib/python3.6/site-packages/transformers/benchmark/benchmark.pyframework_version8 s z"PyTorchBenchmark.framework_version)
model_name
batch_sizesequence_lengthreturnc C s | j |||}| j|S )N)_prepare_inference_func_measure_speed)r r r r
_inferencer r r _inference_speed< s z!PyTorchBenchmark._inference_speedc C s | j |||}| j|S )N)r _measure_memory)r r r r r r r r _inference_memory@ s z"PyTorchBenchmark._inference_memoryc C s | j |||}| j|S )N)_prepare_train_funcr )r r r r _trainr r r _train_speedF s z PyTorchBenchmark._train_speedc C s | j |||}| j|S )N)r$ r" )r r r r r% r r r
_train_memoryJ s z PyTorchBenchmark._train_memoryNc s | j | }| jjrd|_t|do:t|jto:t|jdk}| jj r|ry.|jd }t d|gd}t
||}||} W q t k
r t | dY qX nt |j
|} | j | j| jj t|dr|jn|jj}
tj|
||ftj| jjd| jjr$tjd | jjs td
| j | jjrRtj tj j | W d Q R X n| fd d } fd
d} |j r~| n| }
|
S )NT
architecturesr transformers)fromlistz does not exist. If you just want to test the pretrained model, you might want to set `--only_pretrain_model` or `args.only_pretrain_model=True`.
vocab_size)dtypedevicez&Running training in Mixed Precision...z)Mixed precision is possible only for GPU.c s$ t j d} W d Q R X | S )N)decoder_input_ids)r no_grad)outputs)inference_model input_idsr r encoder_decoder_forward} s
zIPyTorchBenchmark._prepare_inference_func.
.encoder_decoder_forwardc s t j } W d Q R X | S )N)r r/ )r0 )r1 r2 r r encoder_forward s
zAPyTorchBenchmark._prepare_inference_func..encoder_forward) config_dictargs torchscripthasattr
isinstancer( listlenonly_pretrain_model
__import__getattr ImportErrorr __class__evaltor- r+ encoderr randintlongfp16loggerinfois_gpu
ValueErrorhalfr/ Zjittraceis_encoder_decoder)r r r r confighas_model_class_in_config model_classtransformers_module model_clsmodelr+ r3 r4 Z_forwardr )r1 r2 r r P sB
z(PyTorchBenchmark._prepare_inference_funcc s\ | j | }t|do,t|jto,t|jdk}| jj r|ry.|jd }td|gd}t ||}||} W q t
k
r t
| dY qX nt |j |} | jj
rtdn| | j | j| jj t|dr|jn|jj}
tj|
||ftj| jjd | jjr,tjd | jjs$td
| j fd d } fd
d} |j rT| n| }
|
S )Nr( r r) )r* z does not exist. If you just want to test the pretrained model, you might want to set `--only_pretrain_model` or `args.only_pretrain_model=True`.z5Training for torchscript is currently not implementedr+ )r, r- z&Running training in Mixed Precision...z)Mixed precision is possible only for GPU.c s dd } | j | S )N)labelsr )backward)loss)r2 train_modelr r !compute_loss_and_backprob_encoder s zOPyTorchBenchmark._prepare_train_func..compute_loss_and_backprob_encoderc s dd } | j | S )N)r. rT r )rU )rV )r2 rW r r )compute_loss_and_backprob_encoder_decoder s zWPyTorchBenchmark._prepare_train_func..compute_loss_and_backprob_encoder_decoder) r5 r8 r9 r( r: r; r6 r< r= r> r? r r@ r7 NotImplementedErrorZtrainrB r- r+ rC r rD rE rF rG rH rI rJ rK rM )r r r r rN rO rP rQ rR rS r+ rX rY r% r )r2 rW r r$ s>
z$PyTorchBenchmark._prepare_train_func)r c C s yt| j js| j jr,tjd tj|ddd tj|| j jdd}| j jrj| j jrjdd lj j
}| j |j t
|d S tk
r } z| j d| d S d }~X nX d S )
NzRDo inference on TPU or torchscript. Running model 5 times to stabilize compilationr )repeatnumber
r g $@zDoesn't fit on GPU. zN/A)r6 is_tpur7 rG rH timeitr\ Ztorch_xla_tpu_print_metricsZtorch_xla.debug.metricsdebugZmetricsprint_fnZmetrics_reportmin RuntimeError)r funcZruntimesZmeter r r r s$
z PyTorchBenchmark._measure_speed)re r c
C s y| j jrtd}| j jr$tdn| j jrt sBtjd d}qtj d t
j | t
j | j j
}t
j|}|j}t|}t
j n t|}t|trt|n|}| j jrt|}nd }||fS tk
r } z| jd| dS d } ~ X nX d S )Nr) zMemory Benchmarking is currently not implemented for TPU. Please disable memory benchmarking with `--no-memory` or `args.memory=False`zypy3nvml not installed, we won't log GPU memory usage. Install py3nvml (pip install py3nvml) to log information about GPU.N/AzlMeasuring total GPU usage on GPU device. Make sure to not have additional processes running on the same GPU.zDoesn't fit on GPU. )rg N)r6 Ztrace_memory_line_by_liner r_ rZ rI r rG warningrH nvmlZnvmlInitZnvmlDeviceGetHandleByIndexZ
device_idxZnvmlDeviceGetMemoryInfousedr
Z nvmlShutdownr r9 intr rd rb )
r re rL ZmemoryhandleZmeminfoZmax_bytes_in_useZ memory_bytessummaryrf r r r r" s:
z PyTorchBenchmark._measure_memory)__name__
__module__ __qualname__r r6 r Zconfigs frameworkstrpropertyr rk floatr! r
r r r# r& r' r r r$ r r" r r r r r 2 s
: ; r ) __doc__r` typingr r Zconfiguration_utilsr Zmodels.auto.modeling_autor r utilsr r r
Zbenchmark_utilsr r
r r r r r Zbenchmark_argsr Zpy3nvml.py3nvmlZpy3nvmlri
get_loggerrn rG r r r r r s