| 188 | |
| 189 | |
| 190 | def train(args, model, device, train_loader, optimizer, epoch): |
| 191 | model.train() |
| 192 | |
| 193 | # we aren't using `TripletLoss` as the MNIST dataset is simple, so `BCELoss` can do the trick. |
| 194 | criterion = nn.BCELoss() |
| 195 | |
| 196 | for batch_idx, (images_1, images_2, targets) in enumerate(train_loader): |
| 197 | images_1, images_2, targets = images_1.to(device), images_2.to(device), targets.to(device) |
| 198 | optimizer.zero_grad() |
| 199 | outputs = model(images_1, images_2).squeeze() |
| 200 | loss = criterion(outputs, targets) |
| 201 | loss.backward() |
| 202 | optimizer.step() |
| 203 | if batch_idx % args.log_interval == 0: |
| 204 | print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format( |
| 205 | epoch, batch_idx * len(images_1), len(train_loader.dataset), |
| 206 | 100. * batch_idx / len(train_loader), loss.item())) |
| 207 | if args.dry_run: |
| 208 | break |
| 209 | |
| 210 | |
| 211 | def test(model, device, test_loader): |