Testing fastai's lr_find with accelerate.

accelerate_lr_find_1.py

from fastai.vision.all import *
from fastai.distributed import *
from fastai.vision.models.xresnet import *

path = rank0_first(untar_data, URLs.PETS)
def get_y(o): return o[0].isupper()

dls = ImageDataLoaders.from_name_func(
    path=path, 
    fnames=get_image_files(path), 
    valid_pct=0.2, label_func=get_y, 
    item_tfms=Resize(448), 
)

learn = vision_learner(dls, resnet34, metrics=[error_rate, accuracy]).to_fp16()

with learn.distrib_ctx(sync_bn=False, in_notebook=False):
    lrs = learn.lr_find(suggest_funcs=(valley, slide), show_plot=False)
    print(f"{lrs.valley = } | {lrs.slide = }")

accelerate_lr_find_2.py

from fastai.vision.all import *
from fastai.distributed import *
from fastai.vision.models.xresnet import *

from accelerate import Accelerator
accelerate = Accelerator()

path = rank0_first(untar_data, URLs.PETS)
def get_y(o): return o[0].isupper()

dls = ImageDataLoaders.from_name_func(
    path=path, 
    fnames=get_image_files(path), 
    valid_pct=0.2, label_func=get_y, 
    item_tfms=Resize(448), 
)

learn = vision_learner(dls, resnet34, metrics=[error_rate, accuracy]).to_fp16()

with learn.distrib_ctx(sync_bn=False, in_notebook=False):
    lrs = learn.lr_find(suggest_funcs=(valley, slide), show_plot=False)
    print(f"{lrs.valley = } | {lrs.slide = }")

    accelerate.wait_for_everyone()