Skip to content
GitLab
Menu
Projects
Groups
Snippets
Loading...
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in / Register
Toggle navigation
Menu
Open sidebar
OpenDAS
nni
Commits
c037a7c1
Unverified
Commit
c037a7c1
authored
Nov 10, 2019
by
SparkSnail
Committed by
GitHub
Nov 10, 2019
Browse files
Merge pull request #213 from microsoft/master
merge master
parents
49972952
901012eb
Changes
55
Show whitespace changes
Inline
Side-by-side
Showing
20 changed files
with
459 additions
and
170 deletions
+459
-170
src/nni_manager/training_service/remote_machine/gpuScheduler.ts
...i_manager/training_service/remote_machine/gpuScheduler.ts
+22
-1
src/sdk/pycli/setup.py
src/sdk/pycli/setup.py
+10
-10
src/sdk/pynni/nni/assessor.py
src/sdk/pynni/nni/assessor.py
+7
-7
src/sdk/pynni/nni/compression/tensorflow/builtin_pruners.py
src/sdk/pynni/nni/compression/tensorflow/builtin_pruners.py
+10
-6
src/sdk/pynni/nni/compression/tensorflow/builtin_quantizers.py
...dk/pynni/nni/compression/tensorflow/builtin_quantizers.py
+6
-6
src/sdk/pynni/nni/compression/tensorflow/compressor.py
src/sdk/pynni/nni/compression/tensorflow/compressor.py
+114
-59
src/sdk/pynni/nni/compression/torch/builtin_pruners.py
src/sdk/pynni/nni/compression/torch/builtin_pruners.py
+15
-13
src/sdk/pynni/nni/compression/torch/builtin_quantizers.py
src/sdk/pynni/nni/compression/torch/builtin_quantizers.py
+6
-6
src/sdk/pynni/nni/compression/torch/compressor.py
src/sdk/pynni/nni/compression/torch/compressor.py
+152
-35
src/sdk/pynni/nni/evolution_tuner/evolution_tuner.py
src/sdk/pynni/nni/evolution_tuner/evolution_tuner.py
+2
-2
src/sdk/pynni/nni/nas_utils.py
src/sdk/pynni/nni/nas_utils.py
+2
-0
src/sdk/pynni/nni/parameter_expressions.py
src/sdk/pynni/nni/parameter_expressions.py
+1
-1
src/sdk/pynni/nni/ppo_tuner/distri.py
src/sdk/pynni/nni/ppo_tuner/distri.py
+3
-3
src/sdk/pynni/nni/ppo_tuner/policy.py
src/sdk/pynni/nni/ppo_tuner/policy.py
+1
-1
src/sdk/pynni/nni/ppo_tuner/ppo_tuner.py
src/sdk/pynni/nni/ppo_tuner/ppo_tuner.py
+8
-7
src/sdk/pynni/nni/ppo_tuner/util.py
src/sdk/pynni/nni/ppo_tuner/util.py
+1
-1
src/sdk/pynni/nni/smac_tuner/__init__.py
src/sdk/pynni/nni/smac_tuner/__init__.py
+1
-1
src/sdk/pynni/nni/smac_tuner/smac_tuner.py
src/sdk/pynni/nni/smac_tuner/smac_tuner.py
+0
-1
src/sdk/pynni/nni/tuner.py
src/sdk/pynni/nni/tuner.py
+10
-10
src/sdk/pynni/tests/assets/search_space.json
src/sdk/pynni/tests/assets/search_space.json
+88
-0
No files found.
src/nni_manager/training_service/remote_machine/gpuScheduler.ts
View file @
c037a7c1
...
...
@@ -28,6 +28,8 @@ import {
parseGpuIndices
,
RemoteMachineMeta
,
RemoteMachineScheduleResult
,
RemoteMachineTrialJobDetail
,
ScheduleResultType
,
SSHClientManager
}
from
'
./remoteMachineData
'
;
type
SCHEDULE_POLICY_NAME
=
'
random
'
|
'
round-robin
'
;
/**
* A simple GPU scheduler implementation
*/
...
...
@@ -35,13 +37,18 @@ export class GPUScheduler {
private
readonly
machineSSHClientMap
:
Map
<
RemoteMachineMeta
,
SSHClientManager
>
;
private
readonly
log
:
Logger
=
getLogger
();
private
readonly
policyName
:
SCHEDULE_POLICY_NAME
=
'
round-robin
'
;
private
roundRobinIndex
:
number
=
0
;
private
configuredRMs
:
RemoteMachineMeta
[]
=
[];
/**
* Constructor
* @param machineSSHClientMap map from remote machine to sshClient
*/
constructor
(
machineSSHClientMap
:
Map
<
RemoteMachineMeta
,
SSHClientManager
>
)
{
assert
(
machineSSHClientMap
.
size
>
0
);
this
.
machineSSHClientMap
=
machineSSHClientMap
;
this
.
configuredRMs
=
Array
.
from
(
machineSSHClientMap
.
keys
());
}
/**
...
...
@@ -189,7 +196,21 @@ export class GPUScheduler {
private
selectMachine
(
rmMetas
:
RemoteMachineMeta
[]):
RemoteMachineMeta
{
assert
(
rmMetas
!==
undefined
&&
rmMetas
.
length
>
0
);
if
(
this
.
policyName
===
'
random
'
)
{
return
randomSelect
(
rmMetas
);
}
else
if
(
this
.
policyName
===
'
round-robin
'
)
{
return
this
.
roundRobinSelect
(
rmMetas
);
}
else
{
throw
new
Error
(
`Unsupported schedule policy:
${
this
.
policyName
}
`
);
}
}
private
roundRobinSelect
(
rmMetas
:
RemoteMachineMeta
[]):
RemoteMachineMeta
{
while
(
!
rmMetas
.
includes
(
this
.
configuredRMs
[
this
.
roundRobinIndex
%
this
.
configuredRMs
.
length
]))
{
this
.
roundRobinIndex
++
;
}
return
this
.
configuredRMs
[
this
.
roundRobinIndex
++
%
this
.
configuredRMs
.
length
];
}
private
selectGPUsForTrial
(
gpuInfos
:
GPUInfo
[],
requiredGPUNum
:
number
):
GPUInfo
[]
{
...
...
src/sdk/pycli/setup.py
View file @
c037a7c1
import
setuptools
setuptools
.
setup
(
name
=
'nnicli'
,
version
=
'999.0.0-developing'
,
packages
=
setuptools
.
find_packages
(),
name
=
'nnicli'
,
version
=
'999.0.0-developing'
,
packages
=
setuptools
.
find_packages
(),
python_requires
=
'>=3.5'
,
install_requires
=
[
python_requires
=
'>=3.5'
,
install_requires
=
[
'requests'
],
author
=
'Microsoft NNI Team'
,
author_email
=
'nni@microsoft.com'
,
description
=
'nnicli for Neural Network Intelligence project'
,
license
=
'MIT'
,
url
=
'https://github.com/Microsoft/nni'
,
author
=
'Microsoft NNI Team'
,
author_email
=
'nni@microsoft.com'
,
description
=
'nnicli for Neural Network Intelligence project'
,
license
=
'MIT'
,
url
=
'https://github.com/Microsoft/nni'
,
)
src/sdk/pynni/nni/assessor.py
View file @
c037a7c1
...
...
@@ -53,14 +53,14 @@ class Assessor(Recoverable):
to tell whether this trial can be early stopped or not.
This is the abstract base class for all assessors.
Early stopping algorithms should
d
eri
ve
this class and override :meth:`assess_trial` method,
Early stopping algorithms should
inh
eri
t
this class and override :meth:`assess_trial` method,
which receives intermediate results from trials and give an assessing result.
If :meth:`assess_trial` returns :obj:`AssessResult.Bad` for a trial,
it hints NNI framework that the trial is likely to result in a poor final accuracy,
and therefore should be killed to save resource.
If an accessor want's to
get
notified when a trial ends, it can also override :meth:`trial_end`.
If an accessor want's to
be
notified when a trial ends, it can also override :meth:`trial_end`.
To write a new assessor, you can reference :class:`~nni.medianstop_assessor.MedianstopAssessor`'s code as an example.
...
...
@@ -77,7 +77,7 @@ class Assessor(Recoverable):
The NNI framework has little guarantee on ``trial_history``.
This method is not guaranteed to be invoked for each time ``trial_history`` get updated.
It is also possible that a trial's history keeps updat
e
ing after receiving a bad result.
It is also possible that a trial's history keeps updating after receiving a bad result.
And if the trial failed and retried, ``trial_history`` may be inconsistent with its previous value.
The only guarantee is that ``trial_history`` is always growing.
...
...
@@ -96,9 +96,9 @@ class Assessor(Recoverable):
Parameters
----------
trial_job_id: str
trial_job_id
: str
Unique identifier of the trial.
trial_history: list
trial_history
: list
Intermediate results of this trial. The element type is decided by trial code.
Returns
...
...
@@ -114,9 +114,9 @@ class Assessor(Recoverable):
Parameters
----------
trial_job_id: str
trial_job_id
: str
Unique identifier of the trial.
success: bool
success
: bool
True if the trial successfully completed; False if failed or terminated.
"""
...
...
src/sdk/pynni/nni/compression/tensorflow/builtin_pruners.py
View file @
c037a7c1
...
...
@@ -8,16 +8,18 @@ _logger = logging.getLogger(__name__)
class
LevelPruner
(
Pruner
):
def
__init__
(
self
,
config_list
):
def
__init__
(
self
,
model
,
config_list
):
"""
config_list: supported keys:
- sparsity
"""
super
().
__init__
(
config_list
)
super
().
__init__
(
model
,
config_list
)
self
.
mask_list
=
{}
self
.
if_init_list
=
{}
def
calc_mask
(
self
,
weight
,
config
,
op_name
,
**
kwargs
):
def
calc_mask
(
self
,
layer
,
config
):
weight
=
layer
.
weight
op_name
=
layer
.
name
if
self
.
if_init_list
.
get
(
op_name
,
True
):
threshold
=
tf
.
contrib
.
distributions
.
percentile
(
tf
.
abs
(
weight
),
config
[
'sparsity'
]
*
100
)
mask
=
tf
.
cast
(
tf
.
math
.
greater
(
tf
.
abs
(
weight
),
threshold
),
weight
.
dtype
)
...
...
@@ -38,7 +40,7 @@ class AGP_Pruner(Pruner):
https://arxiv.org/pdf/1710.01878.pdf
"""
def
__init__
(
self
,
config_list
):
def
__init__
(
self
,
model
,
config_list
):
"""
config_list: supported keys:
- initial_sparsity
...
...
@@ -47,13 +49,15 @@ class AGP_Pruner(Pruner):
- end_epoch: end epoch number stop update mask
- frequency: if you want update every 2 epoch, you can set it 2
"""
super
().
__init__
(
config_list
)
super
().
__init__
(
model
,
config_list
)
self
.
mask_list
=
{}
self
.
if_init_list
=
{}
self
.
now_epoch
=
tf
.
Variable
(
0
)
self
.
assign_handler
=
[]
def
calc_mask
(
self
,
weight
,
config
,
op_name
,
**
kwargs
):
def
calc_mask
(
self
,
layer
,
config
):
weight
=
layer
.
weight
op_name
=
layer
.
name
start_epoch
=
config
.
get
(
'start_epoch'
,
0
)
freq
=
config
.
get
(
'frequency'
,
1
)
if
self
.
now_epoch
>=
start_epoch
and
self
.
if_init_list
.
get
(
op_name
,
True
)
and
(
...
...
src/sdk/pynni/nni/compression/tensorflow/builtin_quantizers.py
View file @
c037a7c1
...
...
@@ -10,8 +10,8 @@ _logger = logging.getLogger(__name__)
class
NaiveQuantizer
(
Quantizer
):
"""quantize weight to 8 bits
"""
def
__init__
(
self
,
config_list
):
super
().
__init__
(
config_list
)
def
__init__
(
self
,
model
,
config_list
):
super
().
__init__
(
model
,
config_list
)
self
.
layer_scale
=
{}
def
quantize_weight
(
self
,
weight
,
config
,
op_name
,
**
kwargs
):
...
...
@@ -27,12 +27,12 @@ class QAT_Quantizer(Quantizer):
Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inference
http://openaccess.thecvf.com/content_cvpr_2018/papers/Jacob_Quantization_and_Training_CVPR_2018_paper.pdf
"""
def
__init__
(
self
,
config_list
):
def
__init__
(
self
,
model
,
config_list
):
"""
config_list: supported keys:
- q_bits
"""
super
().
__init__
(
config_list
)
super
().
__init__
(
model
,
config_list
)
def
quantize_weight
(
self
,
weight
,
config
,
**
kwargs
):
a
=
tf
.
stop_gradient
(
tf
.
reduce_min
(
weight
))
...
...
@@ -52,12 +52,12 @@ class DoReFaQuantizer(Quantizer):
Zhou et al., DoReFa-Net: Training Low Bitwidth Convolutional Neural Networks with Low Bitwidth Gradients
(https://arxiv.org/abs/1606.06160)
"""
def
__init__
(
self
,
config_list
):
def
__init__
(
self
,
model
,
config_list
):
"""
config_list: supported keys:
- q_bits
"""
super
().
__init__
(
config_list
)
super
().
__init__
(
model
,
config_list
)
def
quantize_weight
(
self
,
weight
,
config
,
**
kwargs
):
a
=
tf
.
math
.
tanh
(
weight
)
...
...
src/sdk/pynni/nni/compression/tensorflow/compressor.py
View file @
c037a7c1
...
...
@@ -6,67 +6,85 @@ _logger = logging.getLogger(__name__)
class
LayerInfo
:
def
__init__
(
self
,
op
):
def
__init__
(
self
,
op
,
weight
,
weight_
op
):
self
.
op
=
op
self
.
name
=
op
.
name
self
.
type
=
op
.
type
self
.
weight
=
weight
self
.
weight_op
=
weight_op
class
Compressor
:
"""Abstract base TensorFlow compressor"""
"""
Abstract base TensorFlow compressor
"""
def
__init__
(
self
,
config_list
):
self
.
_bound_model
=
None
self
.
_config_list
=
config_list
def
__init__
(
self
,
model
,
config_list
):
"""
Record necessary info in class members
Parameters
----------
model : pytorch model
the model user wants to compress
config_list : list
the configurations that users specify for compression
"""
self
.
bound_model
=
model
self
.
config_list
=
config_list
self
.
modules_to_compress
=
[]
def
__call__
(
self
,
model
):
"""Compress given graph with algorithm implemented by subclass.
The graph will be editted and returned.
def
compress
(
self
):
"""
self
.
compress
(
model
)
return
model
Compress the model with algorithm implemented by subclass.
def
compress
(
self
,
model
):
"""Compress given graph with algorithm implemented by subclass.
This will edit the graph.
The model will be instrumented and user should never edit it after calling this method.
`self.modules_to_compress` records all the to-be-compressed layers
"""
assert
self
.
_bound_model
is
None
,
"Each NNI compressor instance can only compress one model"
self
.
_bound_model
=
model
self
.
bind_model
(
model
)
for
op
in
model
.
get_operations
():
layer
=
LayerInfo
(
op
)
config
=
self
.
_select_config
(
layer
)
for
op
in
self
.
bound_model
.
get_operations
():
weight_index
=
_detect_weight_index
(
op
)
if
weight_index
is
None
:
_logger
.
warning
(
'Failed to detect weight for layer %s'
,
op
.
name
)
return
weight_op
=
op
.
inputs
[
weight_index
].
op
weight
=
weight_op
.
inputs
[
0
]
layer
=
LayerInfo
(
op
,
weight
,
weight_op
)
config
=
self
.
select_config
(
layer
)
if
config
is
not
None
:
self
.
_instrument_layer
(
layer
,
config
)
self
.
modules_to_compress
.
append
((
layer
,
config
))
return
self
.
bound_model
def
compress_default_graph
(
self
):
"""Compress the default graph with algorithm implemented by subclass.
This will edit the default graph.
def
get_modules_to_compress
(
self
):
"""
self
.
compress
(
tf
.
get_default_graph
())
To obtain all the to-be-compressed layers.
def
bind_model
(
self
,
model
):
"""This method is called when a model is bound to the compressor.
Compressors can optionally overload this method to do model-specific initialization.
It is guaranteed that only one model will be bound to each compressor instance.
Returns
-------
self.modules_to_compress : list
a list of the layers, each of which is a tuple (`layer`, `config`),
`layer` is `LayerInfo`, `config` is a `dict`
"""
return
self
.
modules_to_compress
def
update_epoch
(
self
,
epoch
,
sess
):
"""If user want to update mask every epoch, user can override this method
def
select_config
(
self
,
layer
):
"""
def
step
(
self
,
sess
):
"""If user want to update mask every step, user can override this method
Find the configuration for `layer` by parsing `self.config_list`
Parameters
----------
layer : LayerInfo
one layer
Returns
-------
ret : config or None
the retrieved configuration for this layer, if None, this layer should
not be compressed
"""
def
_instrument_layer
(
self
,
layer
,
config
):
raise
NotImplementedError
()
def
_select_config
(
self
,
layer
):
ret
=
None
for
config
in
self
.
_
config_list
:
for
config
in
self
.
config_list
:
op_types
=
config
.
get
(
'op_types'
)
if
op_types
==
'default'
:
op_types
=
default_layers
.
op_weight_index
.
keys
()
...
...
@@ -79,35 +97,72 @@ class Compressor:
return
None
return
ret
def
update_epoch
(
self
,
epoch
,
sess
):
"""
If user want to update model every epoch, user can override this method.
This method should be called at the beginning of each epoch
Parameters
----------
epoch : num
the current epoch number
"""
def
step
(
self
,
sess
):
"""
If user want to update mask every step, user can override this method
"""
def
_instrument_layer
(
self
,
layer
,
config
):
"""
This method is implemented in the subclasses, i.e., `Pruner` and `Quantizer`
Parameters
----------
layer : LayerInfo
the layer to instrument the compression operation
config : dict
the configuration for compressing this layer
"""
raise
NotImplementedError
()
class
Pruner
(
Compressor
):
"""
Abstract base TensorFlow pruner
"""
def
calc_mask
(
self
,
weight
,
config
,
op
,
op_type
,
op_name
):
"""Pruners should overload this method to provide mask for weight tensors.
def
calc_mask
(
self
,
layer
,
config
):
"""
Pruners should overload this method to provide mask for weight tensors.
The mask must have the same shape and type comparing to the weight.
It will be applied with `multiply()` operation.
This method works as a subgraph which will be inserted into the bound model.
It will be applied with `mul()` operation on the weight.
This method is effectively hooked to `forward()` method of the model.
Parameters
----------
layer : LayerInfo
calculate mask for `layer`'s weight
config : dict
the configuration for generating the mask
"""
raise
NotImplementedError
(
"Pruners must overload calc_mask()"
)
def
_instrument_layer
(
self
,
layer
,
config
):
# it seems the graph editor can only swap edges of nodes or remove all edges from a node
# it cannot remove one edge from a node, nor can it assign a new edge to a node
# we assume there is a proxy operation between the weight and the Conv2D layer
# this is true as long as the weight is `tf.Value`
# not sure what will happen if the weight is calculated from other operations
weight_index
=
_detect_weight_index
(
layer
)
if
weight_index
is
None
:
_logger
.
warning
(
'Failed to detect weight for layer %s'
,
layer
.
name
)
return
weight_op
=
layer
.
op
.
inputs
[
weight_index
].
op
weight
=
weight_op
.
inputs
[
0
]
mask
=
self
.
calc_mask
(
weight
,
config
,
op
=
layer
.
op
,
op_type
=
layer
.
type
,
op_name
=
layer
.
name
)
new_weight
=
weight
*
mask
tf
.
contrib
.
graph_editor
.
swap_outputs
(
weight_op
,
new_weight
.
op
)
"""
Create a wrapper forward function to replace the original one.
Parameters
----------
layer : LayerInfo
the layer to instrument the mask
config : dict
the configuration for generating the mask
"""
mask
=
self
.
calc_mask
(
layer
,
config
)
new_weight
=
layer
.
weight
*
mask
tf
.
contrib
.
graph_editor
.
swap_outputs
(
layer
.
weight_op
,
new_weight
.
op
)
class
Quantizer
(
Compressor
):
...
...
@@ -133,7 +188,7 @@ def _detect_weight_index(layer):
index
=
default_layers
.
op_weight_index
.
get
(
layer
.
type
)
if
index
is
not
None
:
return
index
weight_indices
=
[
i
for
i
,
op
in
enumerate
(
layer
.
op
.
inputs
)
if
op
.
name
.
endswith
(
'Variable/read'
)]
weight_indices
=
[
i
for
i
,
op
in
enumerate
(
layer
.
inputs
)
if
op
.
name
.
endswith
(
'Variable/read'
)]
if
len
(
weight_indices
)
==
1
:
return
weight_indices
[
0
]
return
None
src/sdk/pynni/nni/compression/torch/builtin_pruners.py
View file @
c037a7c1
...
...
@@ -11,16 +11,17 @@ class LevelPruner(Pruner):
"""Prune to an exact pruning level specification
"""
def
__init__
(
self
,
config_list
):
def
__init__
(
self
,
model
,
config_list
):
"""
config_list: supported keys:
- sparsity
"""
super
().
__init__
(
config_list
)
self
.
mask_list
=
{}
super
().
__init__
(
model
,
config_list
)
self
.
if_init_list
=
{}
def
calc_mask
(
self
,
weight
,
config
,
op_name
,
**
kwargs
):
def
calc_mask
(
self
,
layer
,
config
):
weight
=
layer
.
module
.
weight
.
data
op_name
=
layer
.
name
if
self
.
if_init_list
.
get
(
op_name
,
True
):
w_abs
=
weight
.
abs
()
k
=
int
(
weight
.
numel
()
*
config
[
'sparsity'
])
...
...
@@ -28,10 +29,10 @@ class LevelPruner(Pruner):
return
torch
.
ones
(
weight
.
shape
).
type_as
(
weight
)
threshold
=
torch
.
topk
(
w_abs
.
view
(
-
1
),
k
,
largest
=
False
).
values
.
max
()
mask
=
torch
.
gt
(
w_abs
,
threshold
).
type_as
(
weight
)
self
.
mask_
lis
t
.
update
({
op_name
:
mask
})
self
.
mask_
dic
t
.
update
({
op_name
:
mask
})
self
.
if_init_list
.
update
({
op_name
:
False
})
else
:
mask
=
self
.
mask_
lis
t
[
op_name
]
mask
=
self
.
mask_
dic
t
[
op_name
]
return
mask
...
...
@@ -45,7 +46,7 @@ class AGP_Pruner(Pruner):
https://arxiv.org/pdf/1710.01878.pdf
"""
def
__init__
(
self
,
config_list
):
def
__init__
(
self
,
model
,
config_list
):
"""
config_list: supported keys:
- initial_sparsity
...
...
@@ -54,17 +55,18 @@ class AGP_Pruner(Pruner):
- end_epoch: end epoch number stop update mask, you should make sure start_epoch <= end_epoch
- frequency: if you want update every 2 epoch, you can set it 2
"""
super
().
__init__
(
config_list
)
self
.
mask_list
=
{}
super
().
__init__
(
model
,
config_list
)
self
.
now_epoch
=
0
self
.
if_init_list
=
{}
def
calc_mask
(
self
,
weight
,
config
,
op_name
,
**
kwargs
):
def
calc_mask
(
self
,
layer
,
config
):
weight
=
layer
.
module
.
weight
.
data
op_name
=
layer
.
name
start_epoch
=
config
.
get
(
'start_epoch'
,
0
)
freq
=
config
.
get
(
'frequency'
,
1
)
if
self
.
now_epoch
>=
start_epoch
and
self
.
if_init_list
.
get
(
op_name
,
True
)
and
(
self
.
now_epoch
-
start_epoch
)
%
freq
==
0
:
mask
=
self
.
mask_
lis
t
.
get
(
op_name
,
torch
.
ones
(
weight
.
shape
).
type_as
(
weight
))
mask
=
self
.
mask_
dic
t
.
get
(
op_name
,
torch
.
ones
(
weight
.
shape
).
type_as
(
weight
))
target_sparsity
=
self
.
compute_target_sparsity
(
config
)
k
=
int
(
weight
.
numel
()
*
target_sparsity
)
if
k
==
0
or
target_sparsity
>=
1
or
target_sparsity
<=
0
:
...
...
@@ -73,10 +75,10 @@ class AGP_Pruner(Pruner):
w_abs
=
weight
.
abs
()
*
mask
threshold
=
torch
.
topk
(
w_abs
.
view
(
-
1
),
k
,
largest
=
False
).
values
.
max
()
new_mask
=
torch
.
gt
(
w_abs
,
threshold
).
type_as
(
weight
)
self
.
mask_
lis
t
.
update
({
op_name
:
new_mask
})
self
.
mask_
dic
t
.
update
({
op_name
:
new_mask
})
self
.
if_init_list
.
update
({
op_name
:
False
})
else
:
new_mask
=
self
.
mask_
lis
t
.
get
(
op_name
,
torch
.
ones
(
weight
.
shape
).
type_as
(
weight
))
new_mask
=
self
.
mask_
dic
t
.
get
(
op_name
,
torch
.
ones
(
weight
.
shape
).
type_as
(
weight
))
return
new_mask
def
compute_target_sparsity
(
self
,
config
):
...
...
src/sdk/pynni/nni/compression/torch/builtin_quantizers.py
View file @
c037a7c1
...
...
@@ -10,8 +10,8 @@ logger = logging.getLogger(__name__)
class
NaiveQuantizer
(
Quantizer
):
"""quantize weight to 8 bits
"""
def
__init__
(
self
,
config_list
):
super
().
__init__
(
config_list
)
def
__init__
(
self
,
model
,
config_list
):
super
().
__init__
(
model
,
config_list
)
self
.
layer_scale
=
{}
def
quantize_weight
(
self
,
weight
,
config
,
op_name
,
**
kwargs
):
...
...
@@ -27,12 +27,12 @@ class QAT_Quantizer(Quantizer):
Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inference
http://openaccess.thecvf.com/content_cvpr_2018/papers/Jacob_Quantization_and_Training_CVPR_2018_paper.pdf
"""
def
__init__
(
self
,
config_list
):
def
__init__
(
self
,
model
,
config_list
):
"""
config_list: supported keys:
- q_bits
"""
super
().
__init__
(
config_list
)
super
().
__init__
(
model
,
config_list
)
def
quantize_weight
(
self
,
weight
,
config
,
**
kwargs
):
if
config
[
'q_bits'
]
<=
1
:
...
...
@@ -53,12 +53,12 @@ class DoReFaQuantizer(Quantizer):
Zhou et al., DoReFa-Net: Training Low Bitwidth Convolutional Neural Networks with Low Bitwidth Gradients
(https://arxiv.org/abs/1606.06160)
"""
def
__init__
(
self
,
config_list
):
def
__init__
(
self
,
model
,
config_list
):
"""
config_list: supported keys:
- q_bits
"""
super
().
__init__
(
config_list
)
super
().
__init__
(
model
,
config_list
)
def
quantize_weight
(
self
,
weight
,
config
,
**
kwargs
):
out
=
weight
.
tanh
()
...
...
src/sdk/pynni/nni/compression/torch/compressor.py
View file @
c037a7c1
...
...
@@ -15,49 +15,69 @@ class LayerInfo:
class
Compressor
:
"""Abstract base PyTorch compressor"""
"""
Abstract base PyTorch compressor
"""
def
__init__
(
self
,
config_list
):
self
.
_bound_model
=
None
self
.
_config_list
=
config_list
def
__init__
(
self
,
model
,
config_list
):
"""
Record necessary info in class members
Parameters
----------
model : pytorch model
the model user wants to compress
config_list : list
the configurations that users specify for compression
"""
self
.
bound_model
=
model
self
.
config_list
=
config_list
self
.
modules_to_compress
=
[]
def
__call__
(
self
,
model
):
self
.
compress
(
model
)
return
model
def
compress
(
self
):
"""
Compress the model with algorithm implemented by subclass.
def
compress
(
self
,
model
):
"""Compress the model with algorithm implemented by subclass.
The model will be instrumented and user should never edit it after calling this method.
`self.modules_to_compress` records all the to-be-compressed layers
"""
assert
self
.
_bound_model
is
None
,
"Each NNI compressor instance can only compress one model"
self
.
_bound_model
=
model
self
.
bind_model
(
model
)
for
name
,
module
in
model
.
named_modules
():
for
name
,
module
in
self
.
bound_model
.
named_modules
():
layer
=
LayerInfo
(
name
,
module
)
config
=
self
.
_
select_config
(
layer
)
config
=
self
.
select_config
(
layer
)
if
config
is
not
None
:
self
.
_instrument_layer
(
layer
,
config
)
self
.
modules_to_compress
.
append
((
layer
,
config
))
return
self
.
bound_model
def
bind_model
(
self
,
model
):
"""This method is called when a model is bound to the compressor.
Users can optionally overload this method to do model-specific initialization.
It is guaranteed that only one model will be bound to each compressor instance.
def
get_modules_to_compress
(
self
):
"""
To obtain all the to-be-compressed layers.
def
update_epoch
(
self
,
epoch
):
"""if user want to update model every epoch, user can override this method
Returns
-------
self.modules_to_compress : list
a list of the layers, each of which is a tuple (`layer`, `config`),
`layer` is `LayerInfo`, `config` is a `dict`
"""
return
self
.
modules_to_compress
def
step
(
self
):
"""if user want to update model every step, user can override this method
def
select_config
(
self
,
layer
):
"""
Find the configuration for `layer` by parsing `self.config_list`
def
_instrument_layer
(
self
,
layer
,
config
):
raise
NotImplementedError
()
Parameters
----------
layer : LayerInfo
one layer
def
_select_config
(
self
,
layer
):
Returns
-------
ret : config or None
the retrieved configuration for this layer, if None, this layer should
not be compressed
"""
ret
=
None
for
config
in
self
.
_
config_list
:
for
config
in
self
.
config_list
:
config
[
'op_types'
]
=
self
.
_expand_config_op_types
(
config
)
if
layer
.
type
not
in
config
[
'op_types'
]:
continue
...
...
@@ -68,6 +88,35 @@ class Compressor:
return
None
return
ret
def
update_epoch
(
self
,
epoch
):
"""
If user want to update model every epoch, user can override this method.
This method should be called at the beginning of each epoch
Parameters
----------
epoch : num
the current epoch number
"""
def
step
(
self
):
"""
If user want to update model every step, user can override this method
"""
def
_instrument_layer
(
self
,
layer
,
config
):
"""
This method is implemented in the subclasses, i.e., `Pruner` and `Quantizer`
Parameters
----------
layer : LayerInfo
the layer to instrument the compression operation
config : dict
the configuration for compressing this layer
"""
raise
NotImplementedError
()
def
_expand_config_op_types
(
self
,
config
):
if
config
is
None
:
return
[]
...
...
@@ -79,22 +128,50 @@ class Compressor:
expanded_op_types
.
append
(
op_type
)
return
expanded_op_types
class
Pruner
(
Compressor
):
"""
Abstract base PyTorch pruner
Prune to an exact pruning level specification
Attributes
----------
mask_dict : dict
Dictionary for saving masks, `key` should be layer name and
`value` should be a tensor which has the same shape with layer's weight
"""
def
calc_mask
(
self
,
weight
,
config
,
op
,
op_type
,
op_name
):
"""Pruners should overload this method to provide mask for weight tensors.
def
__init__
(
self
,
model
,
config_list
):
super
().
__init__
(
model
,
config_list
)
self
.
mask_dict
=
{}
def
calc_mask
(
self
,
layer
,
config
):
"""
Pruners should overload this method to provide mask for weight tensors.
The mask must have the same shape and type comparing to the weight.
It will be applied with `mul()` operation.
It will be applied with `mul()` operation
on the weight
.
This method is effectively hooked to `forward()` method of the model.
Parameters
----------
layer : LayerInfo
calculate mask for `layer`'s weight
config : dict
the configuration for generating the mask
"""
raise
NotImplementedError
(
"Pruners must overload calc_mask()"
)
def
_instrument_layer
(
self
,
layer
,
config
):
# TODO: support multiple weight tensors
# create a wrapper forward function to replace the original one
"""
Create a wrapper forward function to replace the original one.
Parameters
----------
layer : LayerInfo
the layer to instrument the mask
config : dict
the configuration for generating the mask
"""
assert
layer
.
_forward
is
None
,
'Each model can only be compressed once'
if
not
_check_weight
(
layer
.
module
):
_logger
.
warning
(
'Module %s does not have parameter "weight"'
,
layer
.
name
)
...
...
@@ -104,16 +181,56 @@ class Pruner(Compressor):
def
new_forward
(
*
inputs
):
# apply mask to weight
old_weight
=
layer
.
module
.
weight
.
data
mask
=
self
.
calc_mask
(
old_weight
,
config
,
op
=
layer
.
module
,
op_type
=
layer
.
type
,
op_name
=
layer
.
name
)
mask
=
self
.
calc_mask
(
layer
,
config
)
layer
.
module
.
weight
.
data
=
old_weight
.
mul
(
mask
)
# calculate forward
ret
=
layer
.
_forward
(
*
inputs
)
# recover original weight
layer
.
module
.
weight
.
data
=
old_weight
return
ret
layer
.
module
.
forward
=
new_forward
def
export_model
(
self
,
model_path
,
mask_path
=
None
,
onnx_path
=
None
,
input_shape
=
None
):
"""
Export pruned model weights, masks and onnx model(optional)
Parameters
----------
model_path : str
path to save pruned model state_dict
mask_path : str
(optional) path to save mask dict
onnx_path : str
(optional) path to save onnx model
input_shape : list or tuple
input shape to onnx model
"""
assert
model_path
is
not
None
,
'model_path must be specified'
for
name
,
m
in
self
.
bound_model
.
named_modules
():
mask
=
self
.
mask_dict
.
get
(
name
)
if
mask
is
not
None
:
mask_sum
=
mask
.
sum
().
item
()
mask_num
=
mask
.
numel
()
_logger
.
info
(
'Layer: %s Sparsity: %.2f'
,
name
,
1
-
mask_sum
/
mask_num
)
print
(
'Layer: %s Sparsity: %.2f'
%
(
name
,
1
-
mask_sum
/
mask_num
))
m
.
weight
.
data
=
m
.
weight
.
data
.
mul
(
mask
)
else
:
_logger
.
info
(
'Layer: %s NOT compressed'
,
name
)
print
(
'Layer: %s NOT compressed'
%
name
)
torch
.
save
(
self
.
bound_model
.
state_dict
(),
model_path
)
_logger
.
info
(
'Model state_dict saved to %s'
,
model_path
)
print
(
'Model state_dict saved to %s'
%
model_path
)
if
mask_path
is
not
None
:
torch
.
save
(
self
.
mask_dict
,
mask_path
)
_logger
.
info
(
'Mask dict saved to %s'
,
mask_path
)
print
(
'Mask dict saved to %s'
%
mask_path
)
if
onnx_path
is
not
None
:
assert
input_shape
is
not
None
,
'input_shape must be specified to export onnx model'
# input info needed
input_data
=
torch
.
Tensor
(
*
input_shape
)
torch
.
onnx
.
export
(
self
.
bound_model
,
input_data
,
onnx_path
)
_logger
.
info
(
'Model in onnx with input shape %s saved to %s'
,
input_data
.
shape
,
onnx_path
)
print
(
'Model in onnx with input shape %s saved to %s'
%
(
input_data
.
shape
,
onnx_path
))
class
Quantizer
(
Compressor
):
"""
...
...
src/sdk/pynni/nni/evolution_tuner/evolution_tuner.py
View file @
c037a7c1
...
...
@@ -158,11 +158,11 @@ class EvolutionTuner(Tuner):
EvolutionTuner is tuner using navie evolution algorithm.
"""
def
__init__
(
self
,
optimize_mode
,
population_size
=
32
):
def
__init__
(
self
,
optimize_mode
=
"maximize"
,
population_size
=
32
):
"""
Parameters
----------
optimize_mode : str
optimize_mode : str
, default 'maximize'
population_size : int
initial population size. The larger population size,
the better evolution performance.
...
...
src/sdk/pynni/nni/nas_utils.py
View file @
c037a7c1
...
...
@@ -265,6 +265,8 @@ def convert_nas_search_space(search_space):
param search_space: raw search space
return: the new search space, mutable_layers will be converted into choice
"""
if
not
isinstance
(
search_space
,
dict
):
return
search_space
ret
=
dict
()
for
k
,
v
in
search_space
.
items
():
if
"_type"
not
in
v
:
...
...
src/sdk/pynni/nni/parameter_expressions.py
View file @
c037a7c1
...
...
@@ -48,7 +48,7 @@ def uniform(low, high, random_state):
high: an float that represent an upper bound
random_state: an object of numpy.random.RandomState
'''
assert
high
>
low
,
'Upper bound must be larger than lower bound'
assert
high
>
=
low
,
'Upper bound must be larger than lower bound'
return
random_state
.
uniform
(
low
,
high
)
...
...
src/sdk/pynni/nni/ppo_tuner/distri.py
View file @
c037a7c1
...
...
@@ -143,14 +143,14 @@ class CategoricalPd(Pd):
re_masked_res
=
tf
.
reshape
(
masked_res
,
[
-
1
,
self
.
size
])
u
=
tf
.
random_uniform
(
tf
.
shape
(
re_masked_res
),
dtype
=
self
.
logits
.
dtype
)
return
tf
.
argmax
(
re_masked_res
-
tf
.
log
(
-
tf
.
log
(
u
)),
axis
=-
1
)
return
tf
.
argmax
(
re_masked_res
-
tf
.
log
(
-
1
*
tf
.
log
(
u
)),
axis
=-
1
)
else
:
u
=
tf
.
random_uniform
(
tf
.
shape
(
self
.
logits
),
dtype
=
self
.
logits
.
dtype
)
return
tf
.
argmax
(
self
.
logits
-
tf
.
log
(
-
tf
.
log
(
u
)),
axis
=-
1
)
return
tf
.
argmax
(
self
.
logits
-
tf
.
log
(
-
1
*
tf
.
log
(
u
)),
axis
=-
1
)
@
classmethod
def
fromflat
(
cls
,
flat
):
return
cls
(
flat
)
return
cls
(
flat
)
# pylint: disable=no-value-for-parameter
class
CategoricalPdType
(
PdType
):
"""
...
...
src/sdk/pynni/nni/ppo_tuner/policy.py
View file @
c037a7c1
...
...
@@ -107,7 +107,7 @@ class PolicyWithValue:
def
sample
(
logits
,
mask_npinf
):
new_logits
=
tf
.
math
.
add
(
logits
,
mask_npinf
)
u
=
tf
.
random_uniform
(
tf
.
shape
(
new_logits
),
dtype
=
logits
.
dtype
)
return
tf
.
argmax
(
new_logits
-
tf
.
log
(
-
tf
.
log
(
u
)),
axis
=-
1
)
return
tf
.
argmax
(
new_logits
-
tf
.
log
(
-
1
*
tf
.
log
(
u
)),
axis
=-
1
)
def
neglogp
(
logits
,
x
):
# return tf.nn.sparse_softmax_cross_entropy_with_logits(logits=self.logits, labels=x)
...
...
src/sdk/pynni/nni/ppo_tuner/ppo_tuner.py
View file @
c037a7c1
...
...
@@ -22,11 +22,9 @@ ppo_tuner.py including:
class PPOTuner
"""
import
os
import
copy
import
logging
import
numpy
as
np
import
json_tricks
from
gym
import
spaces
import
nni
...
...
@@ -236,7 +234,8 @@ class PPOModel:
nextnonterminal
=
1.0
-
trials_info
.
dones
[
t
+
1
]
nextvalues
=
trials_info
.
values
[
t
+
1
]
delta
=
mb_rewards
[
t
]
+
self
.
model_config
.
gamma
*
nextvalues
*
nextnonterminal
-
trials_info
.
values
[
t
]
mb_advs
[
t
]
=
lastgaelam
=
delta
+
self
.
model_config
.
gamma
*
self
.
model_config
.
lam
*
nextnonterminal
*
lastgaelam
lastgaelam
=
delta
+
self
.
model_config
.
gamma
*
self
.
model_config
.
lam
*
nextnonterminal
*
lastgaelam
mb_advs
[
t
]
=
lastgaelam
# pylint: disable=unsupported-assignment-operation
mb_returns
=
mb_advs
+
trials_info
.
values
trials_info
.
update_rewards
(
mb_rewards
,
mb_returns
)
...
...
@@ -536,8 +535,10 @@ class PPOTuner(Tuner):
# generate new trials
self
.
trials_result
=
[
None
for
_
in
range
(
self
.
inf_batch_size
)]
mb_obs
,
mb_actions
,
mb_values
,
mb_neglogpacs
,
mb_dones
,
last_values
=
self
.
model
.
inference
(
self
.
inf_batch_size
)
self
.
trials_info
=
TrialsInfo
(
mb_obs
,
mb_actions
,
mb_values
,
mb_neglogpacs
,
mb_dones
,
last_values
,
self
.
inf_batch_size
)
self
.
trials_info
=
TrialsInfo
(
mb_obs
,
mb_actions
,
mb_values
,
mb_neglogpacs
,
mb_dones
,
last_values
,
self
.
inf_batch_size
)
# check credit and submit new trials
for
_
in
range
(
self
.
credit
):
trial_info_idx
,
actions
=
self
.
trials_info
.
get_next
()
...
...
@@ -581,8 +582,8 @@ class PPOTuner(Tuner):
assert
trial_info_idx
is
not
None
# use mean of finished trials as the result of this failed trial
values
=
[
val
for
val
in
self
.
trials_result
if
val
is
not
None
]
logger
.
warning
(
'zql values:
{0}'
.
format
(
values
)
)
self
.
trials_result
[
trial_info_idx
]
=
(
sum
(
values
)
/
len
(
values
))
if
len
(
values
)
>
0
else
0
logger
.
warning
(
'zql values:
%s'
,
values
)
self
.
trials_result
[
trial_info_idx
]
=
(
sum
(
values
)
/
len
(
values
))
if
values
else
0
self
.
finished_trials
+=
1
if
self
.
finished_trials
==
self
.
inf_batch_size
:
self
.
_next_round_inference
()
...
...
src/sdk/pynni/nni/ppo_tuner/util.py
View file @
c037a7c1
...
...
@@ -56,7 +56,7 @@ def seq_to_batch(h, flat=False):
def
lstm
(
xs
,
ms
,
s
,
scope
,
nh
,
init_scale
=
1.0
):
"""lstm cell"""
nbatch
,
nin
=
[
v
.
value
for
v
in
xs
[
0
].
get_shape
()]
_
,
nin
=
[
v
.
value
for
v
in
xs
[
0
].
get_shape
()]
# the first is nbatch
with
tf
.
variable_scope
(
scope
):
wx
=
tf
.
get_variable
(
"wx"
,
[
nin
,
nh
*
4
],
initializer
=
ortho_init
(
init_scale
))
wh
=
tf
.
get_variable
(
"wh"
,
[
nh
,
nh
*
4
],
initializer
=
ortho_init
(
init_scale
))
...
...
src/sdk/pynni/nni/smac_tuner/__init__.py
View file @
c037a7c1
src/sdk/pynni/nni/smac_tuner/smac_tuner.py
View file @
c037a7c1
...
...
@@ -39,7 +39,6 @@ from nni.utils import OptimizeMode, extract_scalar_reward
from
.convert_ss_to_scenario
import
generate_scenario
class
SMACTuner
(
Tuner
):
"""
Parameters
...
...
src/sdk/pynni/nni/tuner.py
View file @
c037a7c1
...
...
@@ -42,7 +42,7 @@ class Tuner(Recoverable):
A new trial will run with this configuration.
This is the abstract base class for all tuners.
Tuning algorithms should
d
eri
ve
this class and override :meth:`update_search_space`, :meth:`receive_trial_result`,
Tuning algorithms should
inh
eri
t
this class and override :meth:`update_search_space`, :meth:`receive_trial_result`,
as well as :meth:`generate_parameters` or :meth:`generate_multiple_parameters`.
After initializing, NNI will first call :meth:`update_search_space` to tell tuner the feasible region,
...
...
@@ -96,9 +96,9 @@ class Tuner(Recoverable):
Parameters
----------
parameter_id: int
parameter_id
: int
Unique identifier for requested hyper-parameters. This will later be used in :meth:`receive_trial_result`.
**kwargs
:
**kwargs
Unstable parameters which should be ignored by normal users.
Returns
...
...
@@ -129,10 +129,10 @@ class Tuner(Recoverable):
Parameters
----------
parameter_id_list: list of int
parameter_id_list
: list of int
Unique identifiers for each set of requested hyper-parameters.
These will later be used in :meth:`receive_trial_result`.
**kwargs
:
**kwargs
Unstable parameters which should be ignored by normal users.
Returns
...
...
@@ -159,13 +159,13 @@ class Tuner(Recoverable):
Parameters
----------
parameter_id: int
parameter_id
: int
Unique identifier of used hyper-parameters, same with :meth:`generate_parameters`.
parameters
Hyper-parameters generated by :meth:`generate_parameters`.
value
Result from trial (the return value of :func:`nni.report_final_result`).
**kwargs
:
**kwargs
Unstable parameters which should be ignored by normal users.
"""
raise
NotImplementedError
(
'Tuner: receive_trial_result not implemented'
)
...
...
@@ -186,11 +186,11 @@ class Tuner(Recoverable):
Parameters
----------
parameter_id: int
parameter_id
: int
Unique identifier for hyper-parameters used by this trial.
success: bool
success
: bool
True if the trial successfully completed; False if failed or terminated.
**kwargs
:
**kwargs
Unstable parameters which should be ignored by normal users.
"""
...
...
src/sdk/pynni/tests/assets/search_space.json
0 → 100644
View file @
c037a7c1
{
"choice_str"
:
{
"_type"
:
"choice"
,
"_value"
:
[
"cat"
,
"dog"
,
"elephant"
,
"cow"
,
"sheep"
,
"panda"
],
"fail"
:
[
"metis"
,
"gp"
]
},
"choice_int"
:
{
"_type"
:
"choice"
,
"_value"
:
[
42
,
43
,
-1
]
},
"choice_mixed"
:
{
"_type"
:
"choice"
,
"_value"
:
[
0.3
,
"cat"
,
1
,
null
],
"fail"
:
[
"metis"
,
"gp"
]
},
"choice_float"
:
{
"_type"
:
"choice"
,
"_value"
:
[
0.3
,
1
,
2.0
]
},
"choice_single"
:
{
"_type"
:
"choice"
,
"_value"
:
[
1
]
},
"randint_ok"
:
{
"_type"
:
"randint"
,
"_value"
:
[
-2
,
3
]
},
"randint_single"
:
{
"_type"
:
"randint"
,
"_value"
:
[
10
,
11
]
},
"randint_fail_equal"
:
{
"_type"
:
"randint"
,
"_value"
:
[
0
,
0
]
},
"uniform_ok"
:
{
"_type"
:
"uniform"
,
"_value"
:
[
-1.0
,
1.5
]
},
"uniform_equal"
:
{
"_type"
:
"uniform"
,
"_value"
:
[
99.9
,
99.9
]
},
"quniform_ok"
:
{
"_type"
:
"quniform"
,
"_value"
:
[
0.0
,
10.0
,
2.5
]
},
"quniform_clip"
:
{
"_type"
:
"quniform"
,
"_value"
:
[
2.0
,
10.0
,
5.0
]
},
"quniform_clip_2"
:
{
"_type"
:
"quniform"
,
"_value"
:
[
-5.5
,
-0.5
,
6
]
},
"loguniform_ok"
:
{
"_type"
:
"loguniform"
,
"_value"
:
[
0.001
,
100
]
},
"loguniform_equal"
:
{
"_type"
:
"loguniform"
,
"_value"
:
[
1
,
1
]
},
"qloguniform_ok"
:
{
"_type"
:
"qloguniform"
,
"_value"
:
[
0.001
,
100
,
1
]
},
"qloguniform_equal"
:
{
"_type"
:
"qloguniform"
,
"_value"
:
[
2
,
2
,
1
]
},
"normal_ok"
:
{
"_type"
:
"normal"
,
"_value"
:
[
-1.0
,
5.0
]
},
"qnormal_ok"
:
{
"_type"
:
"qnormal"
,
"_value"
:
[
-1.5
,
5.0
,
0.1
]
},
"lognormal_ok"
:
{
"_type"
:
"lognormal"
,
"_value"
:
[
-1.0
,
5.0
]
},
"qlognormal_ok"
:
{
"_type"
:
"qlognormal"
,
"_value"
:
[
-1.5
,
5.0
,
0.1
]
}
}
\ No newline at end of file
Prev
1
2
3
Next
Write
Preview
Markdown
is supported
0%
Try again
or
attach a new file
.
Attach a file
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment