hiroshiaraki / snnlibpy Goto Github PK

train_dataloader = torch.utils.data.DataLoader(
    train_dataset, batch_size=1, shuffle=True)

for step, batch in enumerate( train_dataloader ):
    # Get next input sample.

    inputs = {"X": batch["encoded_image"].view( time, 1, 1, 28, 28 )}
    if gpu:
        inputs = {k: v.cuda() for k, v in inputs.items()}
    label = batch["label"]

    # Run the network on the input.
    network.run( inputs=inputs, time=time, input_time_dim=1 )

    # Optionally plot various simulation information.
    if plot:
        image = batch["image"].view( 28, 28 )

        inpt = inputs["X"].view( time, 784 ).sum( 0 ).view( 28, 28 )
        weights1 = conv_conn.w
        _spikes = {
            "X": spikes["X"].get( "s" ).view( time, -1 ),
            "Y": spikes["Y"].get( "s" ).view( time, -1 ),
        }
        _voltages = {"Y": voltages["Y"].get( "v" ).view( time, -1 )}

        inpt_axes, inpt_ims = plot_input(
            image, inpt, label=label, axes=inpt_axes, ims=inpt_ims
        )
        spike_ims, spike_axes = plot_spikes( _spikes, ims=spike_ims, axes=spike_axes )
        weights1_im = plot_conv2d_weights( weights1, im=weights1_im )
        voltage_ims, voltage_axes = plot_voltages(
            _voltages, ims=voltage_ims, axes=voltage_axes
        )

        #plt.ioff()
        #plt.show()

    plt.ioff()
    # plt.pause( 1 )
    plt.show()

    network.reset_state_variables()  # Reset state variables.