1+ #
2+ # Licensed to the Apache Software Foundation (ASF) under one
3+ # or more contributor license agreements. See the NOTICE file
4+ # distributed with this work for additional information
5+ # regarding copyright ownership. The ASF licenses this file
6+ # to you under the Apache License, Version 2.0 (the
7+ # "License"); you may not use this file except in compliance
8+ # with the License. You may obtain a copy of the License at
9+ #
10+ # http://www.apache.org/licenses/LICENSE-2.0
11+ #
12+ # Unless required by applicable law or agreed to in writing, software
13+ # distributed under the License is distributed on an "AS IS" BASIS,
14+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15+ # See the License for the specific language governing permissions and
16+ # limitations under the License.
17+ #
18+
19+ from singa import layer
20+ from singa import model
21+ from singa import tensor
22+ from singa import opt
23+ from singa import device
24+
25+ import numpy as np
26+
27+ np_dtype = {"float16" : np .float16 , "float32" : np .float32 }
28+
29+ singa_dtype = {"float16" : tensor .float16 , "float32" : tensor .float32 }
30+
31+ class CNN (model .Model ):
32+
33+ def __init__ (self , num_classes = 10 , num_channels = 1 ):
34+ super (CNN , self ).__init__ ()
35+ self .num_classes = num_classes
36+ self .input_size = 128
37+ self .dimension = 4
38+ self .conv1 = layer .Conv2d (num_channels , 32 , 3 , padding = 0 , activation = "RELU" )
39+ self .conv2 = layer .Conv2d (32 , 64 , 3 , padding = 0 , activation = "RELU" )
40+ self .conv3 = layer .Conv2d (64 , 64 , 3 , padding = 0 , activation = "RELU" )
41+ self .linear1 = layer .Linear (128 )
42+ self .linear2 = layer .Linear (num_classes )
43+ self .pooling1 = layer .MaxPool2d (2 , 2 , padding = 0 )
44+ self .pooling2 = layer .MaxPool2d (2 , 2 , padding = 0 )
45+ self .pooling3 = layer .MaxPool2d (2 , 2 , padding = 0 )
46+ self .relu = layer .ReLU ()
47+ self .flatten = layer .Flatten ()
48+ self .softmax_cross_entropy = layer .SoftMaxCrossEntropy ()
49+ self .sigmoid = layer
50+
51+ def forward (self , x ):
52+ y = self .conv1 (x )
53+ y = self .pooling1 (y )
54+ y = self .conv2 (y )
55+ y = self .pooling2 (y )
56+ y = self .conv3 (y )
57+ y = self .pooling3 (y )
58+ y = self .flatten (y )
59+ y = self .linear1 (y )
60+ y = self .relu (y )
61+ y = self .linear2 (y )
62+ return y
63+
64+ def train_one_batch (self , x , y , dist_option , spars ):
65+ out = self .forward (x )
66+ loss = self .softmax_cross_entropy (out , y )
67+
68+ if dist_option == 'plain' :
69+ self .optimizer (loss )
70+ elif dist_option == 'half' :
71+ self .optimizer .backward_and_update_half (loss )
72+ elif dist_option == 'partialUpdate' :
73+ self .optimizer .backward_and_partial_update (loss )
74+ elif dist_option == 'sparseTopK' :
75+ self .optimizer .backward_and_sparse_update (loss ,
76+ topK = True ,
77+ spars = spars )
78+ elif dist_option == 'sparseThreshold' :
79+ self .optimizer .backward_and_sparse_update (loss ,
80+ topK = False ,
81+ spars = spars )
82+ return out , loss
83+
84+ def set_optimizer (self , optimizer ):
85+ self .optimizer = optimizer
86+
87+
88+ class MLP (model .Model ):
89+
90+ def __init__ (self , perceptron_size = 100 , num_classes = 10 ):
91+ super (MLP , self ).__init__ ()
92+ self .num_classes = num_classes
93+ self .dimension = 2
94+
95+ self .relu = layer .ReLU ()
96+ self .linear1 = layer .Linear (perceptron_size )
97+ self .linear2 = layer .Linear (num_classes )
98+ self .softmax_cross_entropy = layer .SoftMaxCrossEntropy ()
99+
100+ def forward (self , inputs ):
101+ y = self .linear1 (inputs )
102+ y = self .relu (y )
103+ y = self .linear2 (y )
104+ return y
105+
106+ def train_one_batch (self , x , y , dist_option , spars ):
107+ out = self .forward (x )
108+ loss = self .softmax_cross_entropy (out , y )
109+
110+ if dist_option == 'plain' :
111+ self .optimizer (loss )
112+ elif dist_option == 'half' :
113+ self .optimizer .backward_and_update_half (loss )
114+ elif dist_option == 'partialUpdate' :
115+ self .optimizer .backward_and_partial_update (loss )
116+ elif dist_option == 'sparseTopK' :
117+ self .optimizer .backward_and_sparse_update (loss ,
118+ topK = True ,
119+ spars = spars )
120+ elif dist_option == 'sparseThreshold' :
121+ self .optimizer .backward_and_sparse_update (loss ,
122+ topK = False ,
123+ spars = spars )
124+ return out , loss
125+
126+ def set_optimizer (self , optimizer ):
127+ self .optimizer = optimizer
128+
129+
130+ def create_model (model_option = 'cnn' , ** kwargs ):
131+ """Constructs a CNN model.
132+
133+ Args:
134+ pretrained (bool): If True, returns a pre-trained model.
135+
136+ Returns:
137+ The created CNN model.
138+ """
139+ model = CNN (** kwargs )
140+ if model_option == 'mlp' :
141+ model = MLP (** kwargs )
142+
143+ return model
144+
145+
146+ __all__ = ['CNN' , 'MLP' , 'create_model' ]
0 commit comments