PyTorch to JAX 移行ガイド(GPUでのCNN学習 | BatchNorm編)

PyTorch to JAX 移行ガイド(GPUでのCNN学習 | BatchNorm編)

Python deeplearning PyTorch # jax tech 背景 前回: PyTorch to JAX 移行ガイド(MLP学習編) JAXベースのNNライブラリであるFlaxを用いて、PyTorchのコードをJAXに移行する方法を紹介しています。特に今回はGPUを用いたCNNの学習を取り上げ、FlaxにおけるBatchNormの使い方について学びます。 python==3.7, jax==0.2.20, flax==0.3.5, optax==0.0.9で動作確認をしています。いずれもバージョンアップとともに実装が変わる可能性もあるので注意 例によってコードはgistにアップロードしてあります。 サンプル: CNNをCIFAR-10で学習する まずはPyTorchから。基本は前回と同じで、JAXへの移行を簡単にするために学習・推論の1ステップをstep関数に切り出しています。あとは、GPUで学習するので.cudaをつけています。ネットワークはよくあるConv-BN-Reluを積み重ねたCNNです。データローダには、前回高速であることが分かったTFD

3 mentions: @jellied_unagi
Date:

Referring Tweets

@jellied_unagi
@jellied_unagi PyTorch to JAX 移行ガイド(GPUでのCNN学習 | BatchNorm編)|jellied_unagi t.co/YfavDNNLEP #zenn 書きました。CIFAR10でのCNNの学習について、PyTorchでは32秒/epochだったのが、JAXでは13秒くらいまで速くなりました。FlaxでBatchNormを使う方法についても詳しく紹介しています。

Bookmark Comments

id:misshiki

Related Entries

Reinforcement Learning Day 2021 - Microsoft Research
Read more Reinforcement Learning Day 2021 - Microsoft Research
1 users, 4 mentions 2021/01/15 21:51
GitHub - robvanvolt/DALLE-datasets: This is a summary of easily available datasets for generalized D...
Read more GitHub - robvanvolt/DALLE-datasets: This is a summary of easily available datasets for generalized D...
3 users, 1 mentions 2021/05/24 00:18
Data Science without Code | Obviously AI
Read more Data Science without Code | Obviously AI
2 users, 6 mentions 2021/07/07 09:09
publications/2021_acm_recsys_transformers4rec at main · NVIDIA-Merlin/publications · GitHub
Read more publications/2021_acm_recsys_transformers4rec at main · NVIDIA-Merlin/publications · GitHub
1 users, 1 mentions 2021/08/02 06:08
元大学教員の高校教諭が総務省「情報II」教材を利用した機械学習の授業案を考えてみた|第1回:総務省教材解説|株式会社ヒューマノーム研究所【AI部】|note
Read more 元大学教員の高校教諭が総務省「情報II」教材を利用した機械学習の授業案を考えてみた|第1回:総務省教材解説|株式会社ヒューマノーム研究所【AI部】|note
1 users, 1 mentions 2022/01/11 07:37

ML-Newsについて

機械学習の技術に関する情報は流速も早いし、分野も多様でキャッチアップが大変です。Twitterで機械学習用のリストを作っても、普段は機械学習以外の話題が多く流れており、効率的に情報収集するのは困難です。

ML-NewsはSNSを情報源とした機械学習に特化したニュースサイトです。機械学習に関する論文ブログライブラリコンペティション発表資料勉強会などの最新の情報を効率的に収集できます。

機械学習を応用した自然言語処理、画像認識、情報検索などの分野の情報や機械学習で必要になるデータ基盤やMLOpsの話題もカバーしています。
安定したサイト運営のためにGitHub sponsorを募集しています。

お知らせ

  • 2021/12/31: デザインを刷新しました
  • 2021/04/08: 日本語Kaggleのカテゴリを新設しました