| |
| |
|
|
|
|
|
|
|
|
|
|
| from caffe2.python import brew, workspace |
| from caffe2.python.model_helper import ModelHelper |
| from caffe2.proto import caffe2_pb2 |
| import logging |
|
|
|
|
| class CNNModelHelper(ModelHelper): |
| """A helper model so we can write CNN models more easily, without having to |
| manually define parameter initializations and operators separately. |
| """ |
|
|
| def __init__(self, order="NCHW", name=None, |
| use_cudnn=True, cudnn_exhaustive_search=False, |
| ws_nbytes_limit=None, init_params=True, |
| skip_sparse_optim=False, |
| param_model=None): |
| logging.warning( |
| "[====DEPRECATE WARNING====]: you are creating an " |
| "object from CNNModelHelper class which will be deprecated soon. " |
| "Please use ModelHelper object with brew module. For more " |
| "information, please refer to caffe2.ai and python/brew.py, " |
| "python/brew_test.py for more information." |
| ) |
|
|
| cnn_arg_scope = { |
| 'order': order, |
| 'use_cudnn': use_cudnn, |
| 'cudnn_exhaustive_search': cudnn_exhaustive_search, |
| } |
| if ws_nbytes_limit: |
| cnn_arg_scope['ws_nbytes_limit'] = ws_nbytes_limit |
| super(CNNModelHelper, self).__init__( |
| skip_sparse_optim=skip_sparse_optim, |
| name="CNN" if name is None else name, |
| init_params=init_params, |
| param_model=param_model, |
| arg_scope=cnn_arg_scope, |
| ) |
|
|
| self.order = order |
| self.use_cudnn = use_cudnn |
| self.cudnn_exhaustive_search = cudnn_exhaustive_search |
| self.ws_nbytes_limit = ws_nbytes_limit |
| if self.order != "NHWC" and self.order != "NCHW": |
| raise ValueError( |
| "Cannot understand the CNN storage order %s." % self.order |
| ) |
|
|
| def ImageInput(self, blob_in, blob_out, use_gpu_transform=False, **kwargs): |
| return brew.image_input( |
| self, |
| blob_in, |
| blob_out, |
| order=self.order, |
| use_gpu_transform=use_gpu_transform, |
| **kwargs |
| ) |
|
|
| def VideoInput(self, blob_in, blob_out, **kwargs): |
| return brew.video_input( |
| self, |
| blob_in, |
| blob_out, |
| **kwargs |
| ) |
|
|
| def PadImage(self, blob_in, blob_out, **kwargs): |
| |
| self.net.PadImage(blob_in, blob_out, **kwargs) |
|
|
| def ConvNd(self, *args, **kwargs): |
| return brew.conv_nd( |
| self, |
| *args, |
| use_cudnn=self.use_cudnn, |
| order=self.order, |
| cudnn_exhaustive_search=self.cudnn_exhaustive_search, |
| ws_nbytes_limit=self.ws_nbytes_limit, |
| **kwargs |
| ) |
|
|
| def Conv(self, *args, **kwargs): |
| return brew.conv( |
| self, |
| *args, |
| use_cudnn=self.use_cudnn, |
| order=self.order, |
| cudnn_exhaustive_search=self.cudnn_exhaustive_search, |
| ws_nbytes_limit=self.ws_nbytes_limit, |
| **kwargs |
| ) |
|
|
| def ConvTranspose(self, *args, **kwargs): |
| return brew.conv_transpose( |
| self, |
| *args, |
| use_cudnn=self.use_cudnn, |
| order=self.order, |
| cudnn_exhaustive_search=self.cudnn_exhaustive_search, |
| ws_nbytes_limit=self.ws_nbytes_limit, |
| **kwargs |
| ) |
|
|
| def GroupConv(self, *args, **kwargs): |
| return brew.group_conv( |
| self, |
| *args, |
| use_cudnn=self.use_cudnn, |
| order=self.order, |
| cudnn_exhaustive_search=self.cudnn_exhaustive_search, |
| ws_nbytes_limit=self.ws_nbytes_limit, |
| **kwargs |
| ) |
|
|
| def GroupConv_Deprecated(self, *args, **kwargs): |
| return brew.group_conv_deprecated( |
| self, |
| *args, |
| use_cudnn=self.use_cudnn, |
| order=self.order, |
| cudnn_exhaustive_search=self.cudnn_exhaustive_search, |
| ws_nbytes_limit=self.ws_nbytes_limit, |
| **kwargs |
| ) |
|
|
| def FC(self, *args, **kwargs): |
| return brew.fc(self, *args, **kwargs) |
|
|
| def PackedFC(self, *args, **kwargs): |
| return brew.packed_fc(self, *args, **kwargs) |
|
|
| def FC_Prune(self, *args, **kwargs): |
| return brew.fc_prune(self, *args, **kwargs) |
|
|
| def FC_Decomp(self, *args, **kwargs): |
| return brew.fc_decomp(self, *args, **kwargs) |
|
|
| def FC_Sparse(self, *args, **kwargs): |
| return brew.fc_sparse(self, *args, **kwargs) |
|
|
| def Dropout(self, *args, **kwargs): |
| return brew.dropout( |
| self, *args, order=self.order, use_cudnn=self.use_cudnn, **kwargs |
| ) |
|
|
| def LRN(self, *args, **kwargs): |
| return brew.lrn( |
| self, *args, order=self.order, use_cudnn=self.use_cudnn, **kwargs |
| ) |
|
|
| def Softmax(self, *args, **kwargs): |
| return brew.softmax(self, *args, use_cudnn=self.use_cudnn, **kwargs) |
|
|
| def SpatialBN(self, *args, **kwargs): |
| return brew.spatial_bn(self, *args, order=self.order, **kwargs) |
|
|
| def SpatialGN(self, *args, **kwargs): |
| return brew.spatial_gn(self, *args, order=self.order, **kwargs) |
|
|
| def InstanceNorm(self, *args, **kwargs): |
| return brew.instance_norm(self, *args, order=self.order, **kwargs) |
|
|
| def Relu(self, *args, **kwargs): |
| return brew.relu( |
| self, *args, order=self.order, use_cudnn=self.use_cudnn, **kwargs |
| ) |
|
|
| def PRelu(self, *args, **kwargs): |
| return brew.prelu(self, *args, **kwargs) |
|
|
| def Concat(self, *args, **kwargs): |
| return brew.concat(self, *args, order=self.order, **kwargs) |
|
|
| def DepthConcat(self, *args, **kwargs): |
| """The old depth concat function - we should move to use concat.""" |
| print("DepthConcat is deprecated. use Concat instead.") |
| return self.Concat(*args, **kwargs) |
|
|
| def Sum(self, *args, **kwargs): |
| return brew.sum(self, *args, **kwargs) |
|
|
| def Transpose(self, *args, **kwargs): |
| return brew.transpose(self, *args, use_cudnn=self.use_cudnn, **kwargs) |
|
|
| def Iter(self, *args, **kwargs): |
| return brew.iter(self, *args, **kwargs) |
|
|
| def Accuracy(self, *args, **kwargs): |
| return brew.accuracy(self, *args, **kwargs) |
|
|
| def MaxPool(self, *args, **kwargs): |
| return brew.max_pool( |
| self, *args, use_cudnn=self.use_cudnn, order=self.order, **kwargs |
| ) |
|
|
| def MaxPoolWithIndex(self, *args, **kwargs): |
| return brew.max_pool_with_index(self, *args, order=self.order, **kwargs) |
|
|
| def AveragePool(self, *args, **kwargs): |
| return brew.average_pool( |
| self, *args, use_cudnn=self.use_cudnn, order=self.order, **kwargs |
| ) |
|
|
| @property |
| def XavierInit(self): |
| return ('XavierFill', {}) |
|
|
| def ConstantInit(self, value): |
| return ('ConstantFill', dict(value=value)) |
|
|
| @property |
| def MSRAInit(self): |
| return ('MSRAFill', {}) |
|
|
| @property |
| def ZeroInit(self): |
| return ('ConstantFill', {}) |
|
|
| def AddWeightDecay(self, weight_decay): |
| return brew.add_weight_decay(self, weight_decay) |
|
|
| @property |
| def CPU(self): |
| device_option = caffe2_pb2.DeviceOption() |
| device_option.device_type = caffe2_pb2.CPU |
| return device_option |
|
|
| @property |
| def GPU(self, gpu_id=0): |
| device_option = caffe2_pb2.DeviceOption() |
| device_option.device_type = workspace.GpuDeviceType |
| device_option.device_id = gpu_id |
| return device_option |
|
|
