Skip to content

Instantly share code, notes, and snippets.

@notwaldorf
Last active June 7, 2023 00:21
Show Gist options
  • Save notwaldorf/7740d4249272fb7a2d509b3ea4240658 to your computer and use it in GitHub Desktop.
Save notwaldorf/7740d4249272fb7a2d509b3ea4240658 to your computer and use it in GitHub Desktop.
ServiceWorker code to cache Tensorflow model shards.

ServiceWorker code to cache TensorFlow model shards.

One of the problems I have when testing giant TensorFlow models in TensorFlow.js is that they're huge (like 500 MB) and they take forever to download, every time I refresh the page. This is how I setup my ServiceWorker code so that at least in testing I only have to download the model once, and then it's saved in the cache for the next time.

<html>
<head>
<script>
// This is just to force HTTPS.
if(location.protocol == 'http:') location.protocol = 'https:';
// Register the Service Worker
if('serviceWorker' in navigator) {
navigator.serviceWorker.register('/sw.js')
.then(reg => console.log('🐳 Service Worker registered', reg))
.catch(err => console.error('⛈ Service Worker **not** registered', err));
}
else {
console.warn('Service Worker not supported in this browser');
}
</script>
</head>
<body>
<p>You're on your own for everything else!</p>
</body>
</html>
{
"name": "your-app-name",
"short_name": "your-app-namet",
"description": "",
"start_url": "/index.html"
}
// Change all of these constants to match your setup.
const CACHE_NAME = 'your-app-name-assets';
// Location of all the shards, i.e. all the files next to your model.json file.
const MODEL_PREFIX = 'https://storage.googleapis.com/your_model_name';
const NUM_SHARDS = 70;
// Make sure your shards match this naming scheme (you might not have an extension, for example)
const SHARDS_NAMING_SCHEME = (i) => `${MODEL_PREFIX}/group1-shard${i}of${NUM_SHARDS}.bin`;
self.addEventListener('install', e => {
e.waitUntil(
(async function() {
const cache = await caches.open(CACHE_NAME);
const resources = [];
resources.push(`${MODEL_PREFIX}/model.json`);
for (let i = 1; i <= NUM_SHARDS; i++) {
resources.push(SHARDS_NAMING_SCHEME(i))
}
// If you have other local files you want to cache, like your index.html or style.css,
// push them onto the resources array before the next line.
await Promise.all([cache.addAll(resources)]);
})()
);
});
self.addEventListener('fetch', e => {
// Fix the trailing slash.
let request = e.request;
if(request.url.endsWith('/')) {
request = new Request(request.url + 'index.html', e);
}
e.respondWith(
caches.match(e.request).then(response => {
if (response) {
console.log('Cached', e.request.url);
return response;
} else {
console.log('Not cached, fetching', e.request.url);
return fetch(e.request);
}
// Or this, if you don't want the console logs for debugging.
// return response || fetch(e.request);
})
);
});
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment