Spaces:
Paused
Paused
File size: 2,145 Bytes
a81a3ae | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 | import os
import pytorch_lightning as pl
from torch.utils.data import DataLoader
from my_dataset import MyDataset
from cldm.logger import ImageLogger
from pytorch_lightning.callbacks import ModelCheckpoint
from cldm.model import create_model, load_state_dict
from cldm.hack import disable_verbosity, enable_sliced_attention
from utils.config import *
# import debugpy; debugpy.listen(('127.0.0.1', 56789)); debugpy.wait_for_client()
if __name__ == '__main__':
os.environ['CUDA_VISIBLE_DEVICES'] = '1' # limit gpu
save_memory = False
disable_verbosity()
if save_memory:
enable_sliced_attention()
# Configs
resume_path = model_root + "control_dresscode_ini.ckpt"
batch_size = 4
logger_freq = 4600 # val
learning_rate = 1.0e-05
sd_locked = True
only_mid_control = False
# First use cpu to load models. Pytorch Lightning will automatically move it to GPUs.
model = create_model('configs/cldm_v2.yaml').cpu()
model.load_state_dict(load_state_dict(resume_path, location='cpu'))
model.learning_rate = learning_rate
model.sd_locked = sd_locked
model.only_mid_control = only_mid_control
# Misc
dataset = MyDataset()
print("******************************************************")
print(len(dataset))
print("******************************************************")
dataloader = DataLoader(dataset, num_workers=0, batch_size=batch_size, shuffle=True)
logger = ImageLogger(batch_frequency=logger_freq)
# ModelCheckpoint
checkpoint_callback = ModelCheckpoint(
monitor=None,
dirpath='./hiera_logs', # dirpath
filename='model_{epoch:02d}-{step:06d}', # file_name
save_top_k=-1, # save all model
save_last=True, # save last model
save_weights_only=False,
mode='min', # Save when the validation indicator is minimized
every_n_train_steps=50000
)
# logger and ModelCheckpoint
callbacks = [logger, checkpoint_callback]
trainer = pl.Trainer(gpus=[1], precision=32, callbacks=callbacks, max_epochs=100)
# Train!
trainer.fit(model, dataloader)
|