(self, x, y, dist_option, spars)
| 49 | return y |
| 50 | |
| 51 | def train_one_batch(self, x, y, dist_option, spars): |
| 52 | out = self.forward(x) |
| 53 | loss = self.softmax_cross_entropy(out, y) |
| 54 | |
| 55 | if dist_option == 'plain': |
| 56 | self.optimizer(loss) |
| 57 | elif dist_option == 'half': |
| 58 | self.optimizer.backward_and_update_half(loss) |
| 59 | elif dist_option == 'partialUpdate': |
| 60 | self.optimizer.backward_and_partial_update(loss) |
| 61 | elif dist_option == 'sparseTopK': |
| 62 | self.optimizer.backward_and_sparse_update(loss, |
| 63 | topK=True, |
| 64 | spars=spars) |
| 65 | elif dist_option == 'sparseThreshold': |
| 66 | self.optimizer.backward_and_sparse_update(loss, |
| 67 | topK=False, |
| 68 | spars=spars) |
| 69 | return out, loss |
| 70 | |
| 71 | def set_optimizer(self, optimizer): |
| 72 | self.optimizer = optimizer |
nothing calls this directly
no test coverage detected