(model, loss, optimizer, inputs, labels)
| 66 | # so we encapsulate it in a function |
| 67 | # Note: inputs and labels are torch tensors |
| 68 | def train(model, loss, optimizer, inputs, labels): |
| 69 | # set the model to training mode |
| 70 | # because batch norm has 2 different modes! |
| 71 | model.train() |
| 72 | |
| 73 | inputs = Variable(inputs, requires_grad=False) |
| 74 | labels = Variable(labels, requires_grad=False) |
| 75 | |
| 76 | # Reset gradient |
| 77 | optimizer.zero_grad() |
| 78 | |
| 79 | # Forward |
| 80 | logits = model.forward(inputs) |
| 81 | output = loss.forward(logits, labels) |
| 82 | |
| 83 | # Backward |
| 84 | output.backward() |
| 85 | |
| 86 | # Update parameters |
| 87 | optimizer.step() |
| 88 | |
| 89 | # what's the difference between backward() and step()? |
| 90 | |
| 91 | return output.item() |
| 92 | |
| 93 | |
| 94 | # similar to train() but not doing the backprop step |
no test coverage detected