Not a member of Pastebin yet? Sign Up, it unlocks many cool features!

1. import tensorflow as tf

2. import numpy as np

3. gpus = tf.config.experimental.list_physical_devices('GPU')

4. if gpus:

5. try:

6. for gpu in gpus:

7. tf.config.experimental.set_memory_growth(gpu, True)

8. except RuntimeError as e:

9. print(e)

11. nclasses = 10

12. nsamples = 3000000

13. bsize = nsamples//20

14. inp_units = 100

15. mod = tf.keras.Sequential([tf.keras.layers.InputLayer(inp_units), tf.keras.layers.Dense(2500, activation='relu'), tf.keras.layers.Dense(2500, activation='relu'), tf.keras.layers.Dense(250, activation='relu'), tf.keras.layers.Dense(nclasses, activation='softmax')])

16. mod.compile(loss='sparse_categorical_crossentropy', optimizer='adam')

18. inpt = np.random.rand(nsamples,inp_units)

19. gtt = np.random.randint(0,nclasses-1,nsamples)

20. dset = tf.data.Dataset.from_tensor_slices((inpt,gtt)).batch(bsize)

22. mod.fit(dset, epochs = 20)