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,
13+ # software distributed under the License is distributed on an
14+ # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+ # KIND, either express or implied. See the License for the
16+ # specific language governing permissions and limitations
17+ # under the License.
18+ #
19+
20+ import numpy as np
21+ from singa import model
22+ from singa import tensor
23+ from singa import layer
24+
25+
26+ np_dtype = {"float16" : np .float16 , "float32" : np .float32 }
27+ singa_dtype = {"float16" : tensor .float16 , "float32" : tensor .float32 }
28+
29+
30+ class MLP (model .Model ):
31+ def __init__ (self , in_features = 10 , perceptron_size = 100 , num_classes = 10 ):
32+ super (MLP , self ).__init__ ()
33+ self .dimension = 2
34+ self .in_features = in_features
35+ self .perceptron_size = perceptron_size
36+ self .num_classes = num_classes
37+ self .relu = layer .ReLU ()
38+ self .linear1 = layer .Linear (self .in_features , self .perceptron_size , bias = True )
39+ self .linear2 = layer .Linear (self .perceptron_size , self .num_classes , bias = True )
40+ self .softmax_cross_entropy = layer .SoftMaxCrossEntropy ()
41+
42+ def forward (self , inputs ):
43+ y = self .linear1 (inputs )
44+ y = self .relu (y )
45+ y = self .linear2 (y )
46+ return y
47+
48+ def train_one_batch (self , x , y , dist_option , spars ):
49+ out = self .forward (x )
50+ loss = self .softmax_cross_entropy (out , y )
51+
52+ if dist_option == 'plain' :
53+ self .optimizer (loss )
54+ elif dist_option == 'half' :
55+ self .optimizer .backward_and_update_half (loss )
56+ elif dist_option == 'partialUpdate' :
57+ self .optimizer .backward_and_partial_update (loss )
58+ elif dist_option == 'sparseTopK' :
59+ self .optimizer .backward_and_sparse_update (loss ,
60+ topK = True ,
61+ spars = spars )
62+ elif dist_option == 'sparseThreshold' :
63+ self .optimizer .backward_and_sparse_update (loss ,
64+ topK = False ,
65+ spars = spars )
66+ return out , loss
67+
68+ def set_optimizer (self , optimizer ):
69+ self .optimizer = optimizer
70+
71+ def create_model (pretrained = False , ** kwargs ):
72+ """Constructs a MLP model.
73+
74+ Args:
75+ pretrained (bool): If True, returns a pre-trained model.
76+
77+ Returns:
78+ The created CNN model.
79+ """
80+ model = MLP (** kwargs )
81+
82+ return model
83+
84+
85+ __all__ = ['MLP' , 'create_model' ]
0 commit comments