def train_step(real_ecg, dim): noise = tf.random.normal(dim) for i in range(disc_steps): with tf.GradientTape() as disc_tape: generated_ecg = generator(noise, training=True) real_output = discriminator(real_ecg, training=True) fake_output = discriminator(generated_ecg, training=True) disc_loss = discriminator_loss(real_output, fake_output) gradients_of_discriminator = disc_tape.gradient(disc_loss, discriminator.trainable_variables) discriminator_optimizer.apply_gradients(zip(gradients_of_discriminator, discriminator.trainable_variables)) ### for tensorboard ### disc_losses.update_state(disc_loss) fake_disc_accuracy.update_state(tf.zeros_like(fake_output), fake_output) real_disc_accuracy.update_state(tf.ones_like(real_output), real_output) ####################### with tf.GradientTape() as gen_tape: generated_ecg = generator(noise, training=True) fake_output = discriminator(generated_ecg, training=True) gen_loss = generator_loss(fake_output) gradients_of_generator = gen_tape.gradient(gen_loss, generator.trainable_variables) generator_optimizer.apply_gradients(zip(gradients_of_generator, generator.trainable_variables)) ### for tensorboard ### gen_losses.update_state(gen_loss) ####################### def train(dataset, epochs, dim): for epoch in tqdm(range(epochs)): for batch in dataset: train_step(batch, dim) disc_losses_list.append(disc_losses.result().numpy()) gen_losses_list.append(gen_losses.result().numpy()) fake_disc_accuracy_list.append(fake_disc_accuracy.result().numpy()) real_disc_accuracy_list.append(real_disc_accuracy.result().numpy()) ### for tensorboard ### # with disc_summary_writer.as_default(): # tf.summary.scalar('loss', disc_losses.result(), step=epoch) # tf.summary.scalar('fake_accuracy', fake_disc_accuracy.result(), step=epoch) # tf.summary.scalar('real_accuracy', real_disc_accuracy.result(), step=epoch) # with gen_summary_writer.as_default(): # tf.summary.scalar('loss', gen_losses.result(), step=epoch) disc_losses.reset_states() gen_losses.reset_states() fake_disc_accuracy.reset_states() real_disc_accuracy.reset_states() ####################### # Save the model every 5 epochs # if (epoch + 1) % 5 == 0: # generate_and_save_ecg(generator, epochs, seed, False) # checkpoint.save(file_prefix = checkpoint_prefix) # Generate after the final epoch display.clear_output(wait=True) generate_and_save_ecg(generator, epochs, seed, False)

下面是将上述代码转换为PyTorch代码的示例：

def train_step(real_ecg, dim):
    noise = torch.randn(dim)
    for i in range(disc_steps):
        generated_ecg = generator(noise, training=True)
        real_output = discriminator(real_ecg, training=True)
        fake_output = discriminator(generated_ecg, training=True)
        disc_loss = discriminator_loss(real_output, fake_output)
        discriminator.zero_grad()
        disc_loss.backward()
        discriminator_optimizer.step()
        
        ### for tensorboard ###
        disc_losses.update(disc_loss)
        fake_disc_accuracy.update(torch.zeros_like(fake_output), fake_output)
        real_disc_accuracy.update(torch.ones_like(real_output), real_output)
        #######################
        
    for i in range(gen_steps):
        generated_ecg = generator(noise, training=True)
        fake_output = discriminator(generated_ecg, training=True)
        gen_loss = generator_loss(fake_output)
        generator.zero_grad()
        gen_loss.backward()
        generator_optimizer.step()
        
        ### for tensorboard ###
        gen_losses.update(gen_loss)
        #######################
        
def train(dataset, epochs, dim):
    for epoch in tqdm(range(epochs)):
        for batch in dataset:
            train_step(batch, dim)
            
        disc_losses_list.append(disc_losses.avg)
        gen_losses_list.append(gen_losses.avg)
        
        fake_disc_accuracy_list.append(fake_disc_accuracy.avg)
        real_disc_accuracy_list.append(real_disc_accuracy.avg)
        
        ### for tensorboard ###
        # with disc_summary_writer.as_default():
        #     tf.summary.scalar('loss', disc_losses.result(), step=epoch)
        #     tf.summary.scalar('fake_accuracy', fake_disc_accuracy.result(), step=epoch)
        #     tf.summary.scalar('real_accuracy', real_disc_accuracy.result(), step=epoch)
        
        # with gen_summary_writer.as_default():
        #     tf.summary.scalar('loss', gen_losses.result(), step=epoch)
        #######################
        
        disc_losses.reset()
        gen_losses.reset()
        
        fake_disc_accuracy.reset()
        real_disc_accuracy.reset()
        
        #######################
        # Save the model every 5 epochs
        # if (epoch + 1) % 5 == 0:
        #     generate_and_save_ecg(generator, epochs, seed, False)
        #     checkpoint.save(file_prefix = checkpoint_prefix)
        
    # Generate after the final epoch
    # display.clear_output(wait=True)
    # generate_and_save_ecg(generator, epochs, seed, False)

注意：上述代码仅作为示例，可能需要根据实际情况进行调整和修改。