examples: set gpus=device_count (#638)

* , gpus=-1

* torch.cuda.device_count()

* .

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Apply suggestions from code review

Co-authored-by: Ethan Harris <ewah1g13@soton.ac.uk>

* .

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Ethan Harris <ewah1g13@soton.ac.uk>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

README.md

@@ -225,6 +225,7 @@ Flash has a [Summarization task](https://lightning-flash.readthedocs.io/en/lates
 
 ```python
 import flash
+import torch
 from flash.core.data.utils import download_data
 from flash.text import SummarizationData, SummarizationTask
 
@@ -244,7 +245,7 @@ datamodule = SummarizationData.from_csv(
 model = SummarizationTask()
 
 # 4. Create the trainer. Run once on data
-trainer = flash.Trainer(max_epochs=1, gpus=1, precision=16)
+trainer = flash.Trainer(max_epochs=1, gpus=torch.cuda.device_count(), precision=16)
 
 # 5. Fine-tune the model
 trainer.finetune(model, datamodule=datamodule)

docs/source/common/finetuning_example.rst

@@ -35,7 +35,7 @@ Here's an example of finetuning.
     model = ImageClassifier(backbone="resnet18", num_classes=datamodule.num_classes)
 
     # 3. Create the trainer (run one epoch for demo)
-    trainer = flash.Trainer(max_epochs=1)
+    trainer = flash.Trainer(max_epochs=1, gpus=torch.cuda.device_count())
 
     # 4. Finetune the model
     trainer.finetune(model, datamodule=datamodule, strategy="freeze")

docs/source/common/training_example.rst

@@ -35,7 +35,7 @@ Here's an example:
     model = ImageClassifier(backbone="resnet18", num_classes=datamodule.num_classes, pretrained=False)
 
     # 3. Create the trainer (run one epoch for demo)
-    trainer = flash.Trainer(max_epochs=1)
+    trainer = flash.Trainer(max_epochs=1, gpus=torch.cuda.device_count())
 
     # 4. Train the model
     trainer.fit(model, datamodule=datamodule)

docs/source/custom_task.rst

@@ -279,7 +279,9 @@ supplying the task itself, and the associated data:
 
     model = RegressionTask(num_inputs=datamodule.train_dataset.num_inputs)
 
-    trainer = flash.Trainer(max_epochs=20, progress_bar_refresh_rate=20, checkpoint_callback=False)
+    trainer = flash.Trainer(
+        max_epochs=20, progress_bar_refresh_rate=20, checkpoint_callback=False, gpus=torch.cuda.device_count()
+    )
     trainer.fit(model, datamodule=datamodule)
 
 

flash_examples/audio_classification.py

@@ -11,6 +11,8 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import torch
+
 import flash
 from flash.audio import AudioClassificationData
 from flash.core.data.utils import download_data
@@ -30,7 +32,7 @@
 model = ImageClassifier(backbone="resnet18", num_classes=datamodule.num_classes)
 
 # 3. Create the trainer and finetune the model
-trainer = flash.Trainer(max_epochs=3)
+trainer = flash.Trainer(max_epochs=3, gpus=torch.cuda.device_count())
 trainer.finetune(model, datamodule=datamodule, strategy=FreezeUnfreeze(unfreeze_epoch=1))
 
 # 4. Predict what's on few images! air_conditioner, children_playing, siren e.t.c

flash_examples/custom_task.py

@@ -157,7 +157,9 @@ class NumpyDataModule(flash.DataModule):
 datamodule = NumpyDataModule.from_numpy(x, y)
 model = RegressionTask(num_inputs=datamodule.train_dataset.num_inputs)
 
-trainer = flash.Trainer(max_epochs=20, progress_bar_refresh_rate=20, checkpoint_callback=False)
+trainer = flash.Trainer(
+    max_epochs=20, progress_bar_refresh_rate=20, checkpoint_callback=False, gpus=torch.cuda.device_count()
+)
 trainer.fit(model, datamodule=datamodule)
 
 predict_data = np.array(

flash_examples/graph_classification.py

@@ -11,6 +11,8 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import torch
+
 import flash
 from flash.core.utilities.imports import _TORCH_GEOMETRIC_AVAILABLE
 from flash.graph import GraphClassificationData, GraphClassifier
@@ -32,7 +34,7 @@
 model = GraphClassifier(num_features=datamodule.num_features, num_classes=datamodule.num_classes)
 
 # 3. Create the trainer and fit the model
-trainer = flash.Trainer(max_epochs=3)
+trainer = flash.Trainer(max_epochs=3, gpus=torch.cuda.device_count())
 trainer.fit(model, datamodule=datamodule)
 
 # 4. Classify some graphs!

flash_examples/image_classification.py

@@ -11,6 +11,8 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import torch
+
 import flash
 from flash.core.data.utils import download_data
 from flash.image import ImageClassificationData, ImageClassifier
@@ -27,7 +29,7 @@
 model = ImageClassifier(backbone="resnet18", num_classes=datamodule.num_classes)
 
 # 3. Create the trainer and finetune the model
-trainer = flash.Trainer(max_epochs=3)
+trainer = flash.Trainer(max_epochs=3, gpus=torch.cuda.device_count())
 trainer.finetune(model, datamodule=datamodule, strategy="freeze")
 
 # 4. Predict what's on a few images! ants or bees?

flash_examples/image_classification_multi_label.py

@@ -11,6 +11,8 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import torch
+
 import flash
 from flash.core.data.utils import download_data
 from flash.image import ImageClassificationData, ImageClassifier
@@ -32,7 +34,7 @@
 model = ImageClassifier(backbone="resnet18", num_classes=datamodule.num_classes, multi_label=True)
 
 # 3. Create the trainer and finetune the model
-trainer = flash.Trainer(max_epochs=3)
+trainer = flash.Trainer(max_epochs=3, gpus=torch.cuda.device_count())
 trainer.finetune(model, datamodule=datamodule, strategy="freeze")
 
 # 4. Predict the genre of a few movies!

flash_examples/integrations/fiftyone/image_classification.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 from itertools import chain
 
+import torch
+
 import flash
 from flash.core.classification import FiftyOneLabels, Labels
 from flash.core.data.utils import download_data
@@ -39,6 +41,7 @@
 )
 trainer = flash.Trainer(
     max_epochs=1,
+    gpus=torch.cuda.device_count(),
     limit_train_batches=1,
     limit_val_batches=1,
 )

flash_examples/integrations/fiftyone/image_classification_fiftyone_datasets.py

@@ -14,6 +14,7 @@
 from itertools import chain
 
 import fiftyone as fo
+import torch
 
 import flash
 from flash.core.classification import FiftyOneLabels, Labels
@@ -53,6 +54,7 @@
 )
 trainer = flash.Trainer(
     max_epochs=1,
+    gpus=torch.cuda.device_count(),
     limit_train_batches=1,
     limit_val_batches=1,
 )

flash_examples/object_detection.py

@@ -11,6 +11,8 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import torch
+
 import flash
 from flash.core.data.utils import download_data
 from flash.image import ObjectDetectionData, ObjectDetector
@@ -23,13 +25,14 @@
     train_folder="data/coco128/images/train2017/",
     train_ann_file="data/coco128/annotations/instances_train2017.json",
     val_split=0.1,
+    batch_size=2,
 )
 
 # 2. Build the task
 model = ObjectDetector(model="retinanet", num_classes=datamodule.num_classes)
 
 # 3. Create the trainer and finetune the model
-trainer = flash.Trainer(max_epochs=3)
+trainer = flash.Trainer(max_epochs=3, gpus=torch.cuda.device_count())
 trainer.finetune(model, datamodule=datamodule)
 
 # 4. Detect objects in a few images!

flash_examples/pointcloud_detection.py

@@ -11,6 +11,8 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import torch
+
 import flash
 from flash.core.data.utils import download_data
 from flash.pointcloud import PointCloudObjectDetector, PointCloudObjectDetectorData
@@ -28,7 +30,9 @@
 model = PointCloudObjectDetector(backbone="pointpillars_kitti", num_classes=datamodule.num_classes)
 
 # 3. Create the trainer and finetune the model
-trainer = flash.Trainer(max_epochs=1, limit_train_batches=1, limit_val_batches=1, num_sanity_val_steps=0)
+trainer = flash.Trainer(
+    max_epochs=1, limit_train_batches=1, limit_val_batches=1, num_sanity_val_steps=0, gpus=torch.cuda.device_count()
+)
 trainer.fit(model, datamodule)
 
 # 4. Predict what's within a few PointClouds?

flash_examples/pointcloud_segmentation.py

@@ -11,6 +11,8 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import torch
+
 import flash
 from flash.core.data.utils import download_data
 from flash.pointcloud import PointCloudSegmentation, PointCloudSegmentationData
@@ -28,7 +30,9 @@
 model = PointCloudSegmentation(backbone="randlanet_semantic_kitti", num_classes=datamodule.num_classes)
 
 # 3. Create the trainer and finetune the model
-trainer = flash.Trainer(max_epochs=1, limit_train_batches=1, limit_val_batches=1, num_sanity_val_steps=0)
+trainer = flash.Trainer(
+    max_epochs=1, limit_train_batches=1, limit_val_batches=1, num_sanity_val_steps=0, gpus=torch.cuda.device_count()
+)
 trainer.fit(model, datamodule)
 
 # 4. Predict what's within a few PointClouds?

flash_examples/semantic_segmentation.py

@@ -11,6 +11,8 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import torch
+
 import flash
 from flash.core.data.utils import download_data
 from flash.image import SemanticSegmentation, SemanticSegmentationData
@@ -39,7 +41,7 @@
 )
 
 # 3. Create the trainer and finetune the model
-trainer = flash.Trainer(max_epochs=3)
+trainer = flash.Trainer(max_epochs=3, gpus=torch.cuda.device_count())
 trainer.finetune(model, datamodule=datamodule, strategy="freeze")
 
 # 4. Segment a few images!

flash_examples/speech_recognition.py

@@ -11,6 +11,8 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import torch
+
 import flash
 from flash.audio import SpeechRecognition, SpeechRecognitionData
 from flash.core.data.utils import download_data
@@ -29,7 +31,7 @@
 model = SpeechRecognition(backbone="facebook/wav2vec2-base-960h")
 
 # 3. Create the trainer and finetune the model
-trainer = flash.Trainer(max_epochs=1)
+trainer = flash.Trainer(max_epochs=1, gpus=torch.cuda.device_count())
 trainer.finetune(model, datamodule=datamodule, strategy="no_freeze")
 
 # 4. Predict on audio files!

flash_examples/style_transfer.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 import os
 
+import torch
+
 import flash
 from flash.core.data.utils import download_data
 from flash.image.style_transfer import StyleTransfer, StyleTransferData
@@ -26,7 +28,7 @@
 model = StyleTransfer(os.path.join(flash.ASSETS_ROOT, "starry_night.jpg"))
 
 # 3. Create the trainer and train the model
-trainer = flash.Trainer(max_epochs=3)
+trainer = flash.Trainer(max_epochs=3, gpus=torch.cuda.device_count())
 trainer.fit(model, datamodule=datamodule)
 
 # 4. Apply style transfer to a few images!

flash_examples/tabular_classification.py

@@ -11,6 +11,8 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import torch
+
 import flash
 from flash.core.data.utils import download_data
 from flash.tabular import TabularClassificationData, TabularClassifier
@@ -30,7 +32,7 @@
 model = TabularClassifier.from_data(datamodule)
 
 # 3. Create the trainer and train the model
-trainer = flash.Trainer(max_epochs=3)
+trainer = flash.Trainer(max_epochs=3, gpus=torch.cuda.device_count())
 trainer.fit(model, datamodule=datamodule)
 
 # 4. Generate predictions from a CSV

flash_examples/template.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import numpy as np
+import torch
 from sklearn import datasets
 
 import flash
@@ -27,7 +28,7 @@
 model = TemplateSKLearnClassifier(num_features=datamodule.num_features, num_classes=datamodule.num_classes)
 
 # 3. Create the trainer and train the model
-trainer = flash.Trainer(max_epochs=3)
+trainer = flash.Trainer(max_epochs=3, gpus=torch.cuda.device_count())
 trainer.fit(model, datamodule=datamodule)
 
 # 4. Classify a few examples

flash_examples/text_classification.py

@@ -11,6 +11,8 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import torch
+
 import flash
 from flash.core.data.utils import download_data
 from flash.text import TextClassificationData, TextClassifier
@@ -30,7 +32,7 @@
 model = TextClassifier(backbone="prajjwal1/bert-medium", num_classes=datamodule.num_classes)
 
 # 3. Create the trainer and finetune the model
-trainer = flash.Trainer(max_epochs=3)
+trainer = flash.Trainer(max_epochs=3, gpus=torch.cuda.device_count())
 trainer.finetune(model, datamodule=datamodule, strategy="freeze")
 
 # 4. Classify a few sentences! How was the movie?

flash_examples/text_classification_multi_label.py

@@ -11,6 +11,8 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import torch
+
 import flash
 from flash.core.data.utils import download_data
 from flash.text import TextClassificationData, TextClassifier
@@ -36,7 +38,7 @@
 )
 
 # 3. Create the trainer and finetune the model
-trainer = flash.Trainer(max_epochs=1)
+trainer = flash.Trainer(max_epochs=1, gpus=torch.cuda.device_count())
 trainer.finetune(model, datamodule=datamodule, strategy="freeze")
 
 # 4. Generate predictions for a few comments!

flash_examples/translation.py

@@ -11,6 +11,8 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import torch
+
 import flash
 from flash.core.data.utils import download_data
 from flash.text import TranslationData, TranslationTask
@@ -30,7 +32,7 @@
 model = TranslationTask(backbone="Helsinki-NLP/opus-mt-en-ro")
 
 # 3. Create the trainer and finetune the model
-trainer = flash.Trainer(max_epochs=3)
+trainer = flash.Trainer(max_epochs=3, gpus=torch.cuda.device_count())
 trainer.finetune(model, datamodule=datamodule)
 
 # 4. Translate something!

flash_examples/video_classification.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 import os
 
+import torch
+
 import flash
 from flash.core.data.utils import download_data
 from flash.video import VideoClassificationData, VideoClassifier
@@ -33,7 +35,7 @@
 model = VideoClassifier(backbone="x3d_xs", num_classes=datamodule.num_classes, pretrained=False)
 
 # 3. Create the trainer and finetune the model
-trainer = flash.Trainer(max_epochs=3)
+trainer = flash.Trainer(max_epochs=3, gpus=torch.cuda.device_count())
 trainer.finetune(model, datamodule=datamodule, strategy="freeze")
 
 # 4. Make a prediction