Comments (7)
hi, how can i enabling float64 precision ?
from neural-tangents.
Sorry for the late reply!
@zhangbububu see https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#double-64bit-precision
@tengandreaxu could you try using Relu(do_stabilize=True)
? https://neural-tangents.readthedocs.io/en/latest/_autosummary/neural_tangents.stax.Relu.html This parameter triggers a way of calculating the nonlinearity kernel in a way that helps prevent numerical overflow.
from neural-tangents.
Thank you so much, Roman. It's no problem at all!
import numpy as np
from neural_tangents import stax
from jax import jit
W_stds = list(range(1, 17))
# W_stds.reverse()
layer_fn = []
for i in range(len(W_stds) - 1):
layer_fn.append(stax.Dense(1, W_std=W_stds[i]))
layer_fn.append(stax.Relu(do_stabilize=True))
layer_fn.append(stax.Dense(1, 1.0, 0.0))
_, _, kernel_fn = stax.serial(*layer_fn)
kernel_fn = jit(kernel_fn, static_argnames="get")
x = np.random.rand(100, 100)
print(kernel_fn(x, x, "ntk"))
results in
[[2.61008562e+20 1.12163820e+20 1.23732785e+20 ... 1.08229372e+20
1.05533967e+20 1.10687273e+20]
[1.12163820e+20 2.92078984e+20 1.31143308e+20 ... 1.16449180e+20
1.15616286e+20 1.19062657e+20]
[1.23732785e+20 1.31143308e+20 3.36093753e+20 ... 1.28641726e+20
1.19473708e+20 1.28997387e+20]
...
[1.08229363e+20 1.16449180e+20 1.28641726e+20 ... 2.74442324e+20
1.07858132e+20 1.20695995e+20]
[1.05533967e+20 1.15616286e+20 1.19473708e+20 ... 1.07858132e+20
2.69344883e+20 1.11830439e+20]
[1.10687273e+20 1.19062657e+20 1.28997387e+20 ... 1.20695995e+20
1.11830439e+20 2.83645061e+20]]
Do you think that there is no sense in having weights drawn from a higher standard deviation as we go deeper into the neural net in the infinite width?
from neural-tangents.
hi, i meet a confuse problem
init_fn, apply_fn, kernel_fn = stax.serial(
stax.Dense(512, W_std=1.5, b_std=0.05), stax.Relu(do_stabilize=True),
stax.Dense(512, W_std=1.5, b_std=0.05), stax.Relu(do_stabilize=True),
stax.Dense(1, W_std=1.5, b_std=0.05)
)
s = 10
l = jnp.pi * -s
r = jnp.pi * s
N_tr = 100
N_te = 5
train_xs = jnp.linspace(l, r , N_tr).reshape(-1, 1).astype(jnp.float64)
train_ys = jnp.sin(train_xs) + jnp.sin(2*train_xs).astype(jnp.float64)
test_xs = jnp.linspace(l, r, N_te).reshape(-1, 1).astype(jnp.float64)
predict_fn = nt.predict.gradient_descent_mse_ensemble(kernel_fn, train_xs,
train_ys, diag_reg=1e-4)
nkt_mean, nkt_covariance = predict_fn(x_test=test_xs, get='ntk',
compute_cov=True)
print(nkt_mean)
if i increate the number of training samples (N_tr), i will get a all NaN nkt_mean
from neural-tangents.
hi, i meet a confuse problem
init_fn, apply_fn, kernel_fn = stax.serial( stax.Dense(512, W_std=1.5, b_std=0.05), stax.Relu(do_stabilize=True), stax.Dense(512, W_std=1.5, b_std=0.05), stax.Relu(do_stabilize=True), stax.Dense(1, W_std=1.5, b_std=0.05) ) s = 10 l = jnp.pi * -s r = jnp.pi * s N_tr = 100 N_te = 5 train_xs = jnp.linspace(l, r , N_tr).reshape(-1, 1).astype(jnp.float64) train_ys = jnp.sin(train_xs) + jnp.sin(2*train_xs).astype(jnp.float64) test_xs = jnp.linspace(l, r, N_te).reshape(-1, 1).astype(jnp.float64) predict_fn = nt.predict.gradient_descent_mse_ensemble(kernel_fn, train_xs, train_ys, diag_reg=1e-4) nkt_mean, nkt_covariance = predict_fn(x_test=test_xs, get='ntk', compute_cov=True) print(nkt_mean)
if i increate the number of training samples (N_tr), i will get a all NaN
nkt_mean
![image](https://private-user-images.githubusercontent.com/41133441/300251255-535d7247-8c23-437a-b83f-6a4843fb618a.png?jwt=eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJnaXRodWIuY29tIiwiYXVkIjoicmF3LmdpdGh1YnVzZXJjb250ZW50LmNvbSIsImtleSI6ImtleTUiLCJleHAiOjE3MjM3ODYxNzQsIm5iZiI6MTcyMzc4NTg3NCwicGF0aCI6Ii80MTEzMzQ0MS8zMDAyNTEyNTUtNTM1ZDcyNDctOGMyMy00MzdhLWI4M2YtNmE0ODQzZmI2MThhLnBuZz9YLUFtei1BbGdvcml0aG09QVdTNC1ITUFDLVNIQTI1NiZYLUFtei1DcmVkZW50aWFsPUFLSUFWQ09EWUxTQTUzUFFLNFpBJTJGMjAyNDA4MTYlMkZ1cy1lYXN0LTElMkZzMyUyRmF3czRfcmVxdWVzdCZYLUFtei1EYXRlPTIwMjQwODE2VDA1MjQzNFomWC1BbXotRXhwaXJlcz0zMDAmWC1BbXotU2lnbmF0dXJlPTYyYTdhNTY2ODFjMDEwZjJlMGQ4NTliZjRlMWExZThiMzQ0OTYxYWZlZGE5MGU2ZDFjMzEzMTI3NWFjN2UxMWMmWC1BbXotU2lnbmVkSGVhZGVycz1ob3N0JmFjdG9yX2lkPTAma2V5X2lkPTAmcmVwb19pZD0wIn0.VhZo5CCK_bqGcqGxpMN2LFLsU-geo7vLE5mKZwQz4N0)
![image](https://private-user-images.githubusercontent.com/41133441/300251267-60b561db-3c21-4bfe-bdd2-5fd922490e5c.png?jwt=eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJnaXRodWIuY29tIiwiYXVkIjoicmF3LmdpdGh1YnVzZXJjb250ZW50LmNvbSIsImtleSI6ImtleTUiLCJleHAiOjE3MjM3ODYxNzQsIm5iZiI6MTcyMzc4NTg3NCwicGF0aCI6Ii80MTEzMzQ0MS8zMDAyNTEyNjctNjBiNTYxZGItM2MyMS00YmZlLWJkZDItNWZkOTIyNDkwZTVjLnBuZz9YLUFtei1BbGdvcml0aG09QVdTNC1ITUFDLVNIQTI1NiZYLUFtei1DcmVkZW50aWFsPUFLSUFWQ09EWUxTQTUzUFFLNFpBJTJGMjAyNDA4MTYlMkZ1cy1lYXN0LTElMkZzMyUyRmF3czRfcmVxdWVzdCZYLUFtei1EYXRlPTIwMjQwODE2VDA1MjQzNFomWC1BbXotRXhwaXJlcz0zMDAmWC1BbXotU2lnbmF0dXJlPWRhZTZmOWRjOTYxYTUyYmJlYjRjOTkwODc2ZTNhMzg5ZTJiOWJmMjJjZWVjMDc5M2NkM2ZiMWY4MzBkMmVhNDcmWC1BbXotU2lnbmVkSGVhZGVycz1ob3N0JmFjdG9yX2lkPTAma2V5X2lkPTAmcmVwb19pZD0wIn0.ItGF10UtjpdMMetpkXDp3OGqzd-FPnBzbYihGdBXGes)
from neural-tangents.
Do you think that there is no sense in having weights drawn from a higher standard deviation as we go deeper into the neural net in the infinite width?
I think so, ideally you would want the mean and variance of your network outputs to match the mean and variance of your training labels, as a sensible prior. But even if your training labels have a large variance, it's common practice to just standardize them (together with test labels) to have mean 0 and variance 1 for best numerical stability.
Then in a Relu network, to have mean zero / variance one outputs (given mean zero, variance one inputs), you would want to set W_std=2**0.5
for all intermediate layers preceding Relus, and W_std=1
for the top layer.
@zhangbububu replied in your separate thread, let's continue there.
from neural-tangents.
Thank you for your prompt help Roman!
from neural-tangents.
Related Issues (20)
- Question regarding OOM issues HOT 3
- Question regarding lr in Neural Tangents Cookbook
- eNTK implementation uses deprecated xla attribute HOT 2
- Colab notebooks issue HOT 2
- How to obtain aleatoric uncertainty? HOT 2
- How to compute the empirical after kernel? HOT 1
- pip install issues HOT 2
- Erf function goes beyond [-1,1] HOT 2
- using stax.Cos(a=1.0, b=1.0, c=0.0) to get kernel from conv layer gives error HOT 2
- NTK is not PD
- stax.serial PSDness HOT 1
- How to use batch to gradient_descent_mse_ensemble ? HOT 1
- NKT_mean output Nan, when the number of training sample is increased HOT 3
- Inefficient jacobian computation for embedding layers. HOT 1
- Question regarding the cookbook
- Calling the empirical kernel function with different parameters returns same result
- Tag a release HOT 2
- Transformer Block Implementation gives NotImplementedError
- High Memory Usage on Infinite-Width NTK for GPU Only
Recommend Projects
-
React
A declarative, efficient, and flexible JavaScript library for building user interfaces.
-
Vue.js
🖖 Vue.js is a progressive, incrementally-adoptable JavaScript framework for building UI on the web.
-
Typescript
TypeScript is a superset of JavaScript that compiles to clean JavaScript output.
-
TensorFlow
An Open Source Machine Learning Framework for Everyone
-
Django
The Web framework for perfectionists with deadlines.
-
Laravel
A PHP framework for web artisans
-
D3
Bring data to life with SVG, Canvas and HTML. 📊📈🎉
-
Recommend Topics
-
javascript
JavaScript (JS) is a lightweight interpreted programming language with first-class functions.
-
web
Some thing interesting about web. New door for the world.
-
server
A server is a program made to process requests and deliver data to clients.
-
Machine learning
Machine learning is a way of modeling and interpreting data that allows a piece of software to respond intelligently.
-
Visualization
Some thing interesting about visualization, use data art
-
Game
Some thing interesting about game, make everyone happy.
Recommend Org
-
Facebook
We are working to build community through open source technology. NB: members must have two-factor auth.
-
Microsoft
Open source projects and samples from Microsoft.
-
Google
Google ❤️ Open Source for everyone.
-
Alibaba
Alibaba Open Source for everyone
-
D3
Data-Driven Documents codes.
-
Tencent
China tencent open source team.
from neural-tangents.