Lightning-AI / Lightning-AI/pytorch-lightning

Proposal for Enhanced State Management and Granularity in Trainer Hooks

Open
#19,215 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

feature trainer
Dominant language
Python
Stars
31.4k
Forks
3.8k
Avg merge
6d 7h
Merged PRs (30d)
6

Description

### Description & Motivation

## Introduction:
Greetings PyTorch Lightning Team,

I've been working with PyTorch Lightning and have encountered some confusion regarding the `stage` parameter in setup, teardown, and other hooks. The parameter values ("fit", "validate", "test", "predict") present some challenges, particularly since validation is typically part of the fitting process. This makes it difficult to distinguish between training and validation stages in certain scenarios, such as when switching between training and validation augmentations.

Additionally, upon exploring the PyTorch Lightning Trainer documentation, I noticed potential inconsistencies in state representation. The `trainer.state.fn` seems to be treated more like the state, while the actual `trainer.state.status` appears underutilized. This inconsistency can be slightly perplexing, especially for new users.

## Issue Description:
The current setup complicates the implementation of certain features, such as dynamically adjusting behaviors for different training stages. For example, determining whether we are in a training or validation phase within the `setup` function based on the `stage` parameter is not straightforward.

## Proposal:
1. **Rename 'stage' to 'state'**: This would reflect that it's a more general object encompassing various aspects of the training process.

2. **Introduce Granularity in State Parameter**: Provide an additional parameter in the trainer that sets the granularity of what is passed as "state". The options could be "state", "stage", "fn", and "status". Selecting "state" would pass a dictionary with keys "fn", "stage", "status" with their corresponding values. Choosing one of the other options would pass only the relevant value, allowing users to decide the level of granularity they require for each training phase.

## Benefits:
- This proposed change would bring more clarity and consistency to the framework, especially regarding the stages of training and validation.
- It allows users to have finer control over the behavior of their models and data handling during different training phases.
- It simplifies understanding for new users, making the framework more approachable.

## Examples:
- A practical use case would be in dynamically selecting data augmentation strategies for training and validation phases within the `setup` method, based on the detailed state information.
- Another scenario could be in resource allocation and optimization, where different stages of the model lifecycle might require varied resources or configurations.

## Request for Feedback:
I would love to hear the community's thoughts on this proposal. Any feedback, suggestions, or insights would be greatly appreciated. I believe these changes could significantly enhance the user experience and functionality of PyTorch Lightning.

Looking forward to the discussion!

Best regards,
AgentDS007

### Pitch

_No response_

### Alternatives

_No response_

### Additional context

_No response_

cc @lantiga @borda @justusschock

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

No files or tests are named. Start by reviewing the Trainer documentation and the setup/teardown hooks alongside trainer.state.fn and trainer.state.status. Done would require an agreed state API with specified behavior for the proposed granularity options.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.