I’ve got something that will work in the meantime @MichaelScofield
def FlatCosAnnealSchedulerWOverSample(learn, lr:float=4e-3, tot_epochs:int=1, moms:Floats=(0.95,0.999),
start_pct:float=0.72, curve='cosine', weights:torch.Tensor=None):
"Manage FCFit trainnig as found in the ImageNette experiments"
ds, dl = learn.data.train_ds, learn.data.train_dl
labels = ds.y.items
assert np.issubdtype(labels.dtype, np.integer), "Can only oversample integer values"
_,label_counts = np.unique(labels,return_counts=True)
if weights is None: weights = torch.DoubleTensor((1/label_counts)[labels])
total_len_oversample = int(learn.data.c*np.max(label_counts))
sampler = WeightedRandomSampler(weights, total_len_oversample)
learn.data.train_dl = dl.new(shuffle=False, sampler=sampler)
n = len(learn.data.train_dl)
anneal_start = int(n * tot_epochs * start_pct)
batch_finish = ((n * tot_epochs) - anneal_start)
if curve=="cosine": curve_type=annealing_cos
elif curve=="linear": curve_type=annealing_linear
elif curve=="exponential": curve_type=annealing_exp
else: raiseValueError(f"annealing type not supported {curve}")
phase0 = TrainingPhase(anneal_start).schedule_hp('lr', lr).schedule_hp('mom', moms[0])
phase1 = TrainingPhase(batch_finish).schedule_hp('lr', lr, anneal=curve_type).schedule_hp('mom', moms[1])
phases = [phase0, phase1]
return GeneralScheduler(learn, phases)
def fit_fcO(learn:Learner, tot_epochs:int=1, lr:float=defaults.lr, moms:Tuple[float,float]=(0.95,0.85), start_pct:float=0.72,
wd:float=None, callbacks:Optional[CallbackList]=None)->None:
"Fit a model with Flat Cosine Annealing"
max_lr = learn.lr_range(lr)
callbacks = listify(callbacks)
callbacks.append(FlatCosAnnealSchedulerWOverSample(learn, lr, moms=moms, start_pct=start_pct, tot_epochs=tot_epochs))
learn.fit(tot_epochs, max_lr, wd=wd, callbacks=callbacks)
To use:
fit_fcO(learn, n_epochs, lr)
I’m unsure quite where it needs to fall in that order bit, I need to do more work on that but this is what I can fit in for tonight 
Thanks for the assist @ilovescience