| 28 | |
| 29 | |
| 30 | class MLP(model.Model): |
| 31 | def __init__(self, data_size=10, perceptron_size=100, num_classes=10): |
| 32 | super(MLP, self).__init__() |
| 33 | self.num_classes = num_classes |
| 34 | self.dimension = 2 |
| 35 | |
| 36 | self.relu = layer.ReLU() |
| 37 | self.linear1 = layer.Linear(perceptron_size) |
| 38 | self.linear2 = layer.Linear(num_classes) |
| 39 | self.softmax_cross_entropy = layer.SoftMaxCrossEntropy() |
| 40 | |
| 41 | def forward(self, inputs): |
| 42 | y = self.linear1(inputs) |
| 43 | y = self.relu(y) |
| 44 | y = self.linear2(y) |
| 45 | return y |
| 46 | |
| 47 | def train_one_batch(self, x, y, dist_option, spars): |
| 48 | out = self.forward(x) |
| 49 | loss = self.softmax_cross_entropy(out, y) |
| 50 | |
| 51 | if dist_option == "plain": |
| 52 | self.optimizer(loss) |
| 53 | elif dist_option == "half": |
| 54 | self.optimizer.backward_and_update_half(loss) |
| 55 | elif dist_option == "partialUpdate": |
| 56 | self.optimizer.backward_and_partial_update(loss) |
| 57 | elif dist_option == "sparseTopK": |
| 58 | self.optimizer.backward_and_sparse_update(loss, topK=True, spars=spars) |
| 59 | elif dist_option == "sparseThreshold": |
| 60 | self.optimizer.backward_and_sparse_update(loss, topK=False, spars=spars) |
| 61 | return out, loss |
| 62 | |
| 63 | def set_optimizer(self, optimizer): |
| 64 | self.optimizer = optimizer |
| 65 | |
| 66 | |
| 67 | def create_model(pretrained=False, **kwargs): |
no outgoing calls
no test coverage detected