3
EgV2 @ s d Z ddlmZmZ ddlmZmZ ddlmZm Z m
Z
m Z ddl Z
G dd deZG dd d eZG d
d d eZdd d
Zdd ZG dd deZdS )aH
Implementation of an in-memory dataset with structured schema.
Use this to store and iterate through datasets with complex schema that
fit in memory.
Iterating through entries of this dataset is very fast since the dataset
is stored as a set of native Caffe2 tensors, thus no type conversion or
deserialization is necessary.
)core workspace)ReaderWriter)Structfrom_blob_listfrom_column_listInitEmptyRecordNc @ s. e Zd Zd ddZdd Zdd Zd d
Zd S )
_DatasetReader Fc C s< t j| |j || _|p"|jd | _|| _|| _d| _dS )z7Don't call this directly. Instead, use dataset.reader()_cursorN)r __init__contentdatasetname
batch_sizeenforce_batch_sizecursor)selfr r r r r A/usr/local/lib64/python3.6/site-packages/caffe2/python/dataset.pyr
s z_DatasetReader.__init__c C s, | j d kr(|jg |j| j| jjd| _ d S )N)fields)r CreateTreeCursorNextScopedBlobr r r )r init_netexit_netr r r setup_ex# s
z_DatasetReader.setup_exc
C sz | j std| jj }tj|j| jF |j| j g|j |j
| j | j d}tj
|}|j|d g|fS Q R X d S )Nzsetup not called.)r r r )r AssertionErrorr r r NameScopeNextNamer Z
ReadNextBatch field_blobs field_namesr r output_to_listIsEmpty)r read_netr r r r r read* s
z_DatasetReader.readc C s |j | jgg d S )N) ResetCursorr )r netr r r reset6 s z_DatasetReader.resetN)r F)__name__
__module__ __qualname__r
r r% r( r r r r r
s
r
c @ s@ e Zd ZdddZdd Zdd Zd d
Zdd d
Zdd Zd S )_DatasetRandomReaderr Fc C sH t j| |j || _d| _|p(|jd | _|| _|| _|| _|| _ dS )z>Don't call this directly. Instead, use dataset.random_reader()Nr )
r r
r r r r indicesr loop_overr )r r r r- r r. r r r r r
; s z _DatasetRandomReader.__init__c C s, | j d kr(|jg |j| j| jjd| _ d S )N)r )r r r r r r )r r r r r r r G s
z _DatasetRandomReader.setup_exc C s |j | jgg d S )N)r& r )r r' r r r r( N s z_DatasetRandomReader.resetc C s2 | j | |j| jg| jj j d}|| _d S )Noffsets)r( Z
ComputeOffsetr r r r r/ )r r' r/ r r r
computeoffsetQ s
z"_DatasetRandomReader.computeoffsetNc C sh | j j }d}|r4||j ks&td|j j|}| j| |j| jg|j d|||d}|| _ d S )Nr zMust be valid field.r- )sort_by_field_idx shuffle_sizer )
r r r! r indexr( ZSortAndShuffler r r- )r r' Z
sort_by_fieldr2 r r r1 r- r r r sort_and_shuffleX s
z%_DatasetRandomReader.sort_and_shufflec C s | j std| js td| js*td| jj }tj|j| j R |j
| j | j| jg|j |j | j
| j| jd}tj|}|j|d g|fS Q R X d S )Nzsetup_ex not calledzsort_and_shuffle not calledzcomputeoffset not called)r r r. r )r r r- r/ r r r r r r ZReadRandomBatchr r! r r r. r" r# )r r$ r r r r r r% k s
z_DatasetRandomReader.read)r FF)Nr r ) r) r* r+ r
r r( r0 r5 r% r r r r r, : s
r, c @ s, e Zd Zdd Zdd Zdd Zdd Zd S )
_DatasetWriterc C s || _ d| _dS )z7Don't call this directly. Use dataset.writer() instead.N)_contentmutex)r r r r r r
} s z_DatasetWriter.__init__c C s | j d kr|jg | _ d S )N)r8 Z CreateMutex)r r r r r r r s
z_DatasetWriter.setup_exc C sz | j dk std| jj }t|t|ksDtdt|t|f |j|g | jj d |j| j g| t| | dS )a
Add operations to `net` that append the blobs in `fields` to the end
of the dataset. An additional operator will also be added that checks
the consistency of the data in `fields` against the dataset schema.
Args:
writer_net: The net that will contain the Append operators.
fields: A list of BlobReference to be appeneded to this dataset.
Nzsetup not called.zExpected %s fields, got %s.)r ) r8 r r7 r lenZCheckDatasetConsistencyr! Z AtomicAppendlist)r Z
writer_netr r r r r write s
z_DatasetWriter.writec C s dS )z+Commit is a no-op for an in-memory dataset.Nr )r Z
finish_netr r r commit s z_DatasetWriter.commitN)r) r* r+ r
r r; r< r r r r r6 | s r6 c C sH t | tjstdtj||d}| j| j|d}tj t
|| |S )a
Create a 'constant' by first creating an external input in the given
net, and then feeding the corresponding blob with its provided value
in the current workspace. The name is automatically generated in order
to avoid clashes with existing blob names.
z net must be a core.Net instance.)dtype)prefix)
isinstancer Netr nparrayZAddExternalInputr r ZFeedBlobstr)r' valuer= r Zblobr r r Const s
rE c C s, t jd}|j|gg t j| ||dddS )N
report_netT )rF Zconcurrent_substepsZreport_interval)r r@ ZPrintZexecution_step)r r ZsubstepsZ rows_readrF r r r execution_step_with_progress s
rH c @ sp e Zd ZdZdddZdd Zdd Zd d
Zd d Zd
d Z dd Z
dd Z d ddZ d ddZ
d ddZdS ) DatasetaS Represents an in-memory dataset with fixed schema.
Use this to store and iterate through datasets with complex schema that
fit in memory.
Iterating through entries of this dataset is very fast since the dataset
is stored as a set of native Caffe2 tensors, thus no type conversion or
deserialization is necessary.
Nc C sl t |ts t |ts tdt |tr.t|}|| _|j | _|j | _|pNd| _ |j
rb|j nd| _ dS )a Create an un-initialized dataset with schema provided by `fields`.
Before this dataset can be used, it must be initialized, either by
`init_empty` or `init_from_dataframe`.
Args:
fields: either a schema.Struct or a list of field names in a format
compatible with the one described in schema.py.
name: optional name to prepend to blobs that will store the data.
z