| 504 | |
| 505 | |
| 506 | def fit(self, X): |
| 507 | d_costs = [] |
| 508 | g_costs = [] |
| 509 | |
| 510 | N = len(X) |
| 511 | n_batches = N // BATCH_SIZE |
| 512 | total_iters = 0 |
| 513 | for i in range(EPOCHS): |
| 514 | print("epoch:", i) |
| 515 | np.random.shuffle(X) |
| 516 | for j in range(n_batches): |
| 517 | t0 = datetime.now() |
| 518 | |
| 519 | if type(X[0]) is str: |
| 520 | # is celeb dataset |
| 521 | batch = util.files2images_theano( |
| 522 | X[j*BATCH_SIZE:(j+1)*BATCH_SIZE] |
| 523 | ) |
| 524 | |
| 525 | else: |
| 526 | # is mnist dataset |
| 527 | batch = X[j*BATCH_SIZE:(j+1)*BATCH_SIZE] |
| 528 | |
| 529 | Z = np.random.uniform(-1, 1, size=(BATCH_SIZE, self.latent_dims)) |
| 530 | |
| 531 | # train the discriminator |
| 532 | d_cost, d_acc = self.train_d(batch, Z) |
| 533 | d_costs.append(d_cost) |
| 534 | |
| 535 | # train the generator |
| 536 | g_cost1 = self.train_g(Z) |
| 537 | g_cost2 = self.train_g(Z) |
| 538 | g_costs.append((g_cost1 + g_cost2)/2) # just use the avg |
| 539 | |
| 540 | print(" batch: %d/%d - dt: %s - d_acc: %.2f" % (j+1, n_batches, datetime.now() - t0, d_acc)) |
| 541 | |
| 542 | # save samples periodically |
| 543 | total_iters += 1 |
| 544 | if total_iters % SAVE_SAMPLE_PERIOD == 0: |
| 545 | print("saving a sample...") |
| 546 | samples = self.sample(64) # shape is (64, D, D, color) |
| 547 | |
| 548 | # for convenience |
| 549 | d = self.img_length |
| 550 | |
| 551 | if samples.shape[-1] == 1: |
| 552 | # if color == 1, we want a 2-D image (N x N) |
| 553 | samples = samples.reshape(64, d, d) |
| 554 | flat_image = np.empty((8*d, 8*d)) |
| 555 | |
| 556 | k = 0 |
| 557 | for i in range(8): |
| 558 | for j in range(8): |
| 559 | flat_image[i*d:(i+1)*d, j*d:(j+1)*d] = samples[k].reshape(d, d) |
| 560 | k += 1 |
| 561 | |
| 562 | # plt.imshow(flat_image, cmap='gray') |
| 563 | else: |