283 lines
11 KiB
Markdown
283 lines
11 KiB
Markdown
# Quick Start
|
||
|
||
This chapter takes stable diffusion v1.5 as an example, demonstrating how to build a neural network architecture from scratch, as well as how to train and test a model based on that structure.
|
||
|
||
# 1. Defining Model Structure with the Network Class
|
||
|
||
The network class encompasses methods for defining, training, and testing the model. First, we initialize a network class and then define the required submodules: autoencoder, UNet, embedder, and loss.
|
||
|
||
|
||
|
||
```python
|
||
@MODELS.register_class()
|
||
class LatentDiffusion(TrainModule):
|
||
def __init__(self, cfg, logger=None):
|
||
super().__init__(cfg, logger=logger)
|
||
self.model_config = cfg.DIFFUSION_MODEL
|
||
self.first_stage_config = cfg.FIRST_STAGE_MODEL
|
||
self.cond_stage_config = cfg.COND_STAGE_MODEL
|
||
self.loss_config = cfg.get('LOSS', None)
|
||
|
||
self.model = BACKBONES.build(self.model_config, logger=self.logger)
|
||
self.first_stage_model = MODELS.build(self.first_stage_config,
|
||
logger=self.logger)
|
||
self.cond_stage_model = EMBEDDERS.build(self.cond_stage_config,
|
||
logger=self.logger)
|
||
if self.loss_config:
|
||
self.loss = LOSSES.build(self.loss_config, logger=self.logger)
|
||
|
||
# Other module definition.
|
||
```
|
||
|
||
# 2. Implementing the Training and Testing Code for the Network class
|
||
|
||
Each network class relies on forward_train and forward_test functions to define its training and testing processes. In SD1.5, forward_train predicts the noise for each timestep t and conducts loss calculation.
|
||
|
||
```python
|
||
def forward_train(self, image=None, noise=None, prompt=None, **kwargs):
|
||
x_start = self.encode_first_stage(image, **kwargs)
|
||
t = torch.randint(0,
|
||
self.num_timesteps, (x_start.shape[0], ),
|
||
device=x_start.device).long()
|
||
context = {}
|
||
if prompt and self.cond_stage_model:
|
||
zeros = (torch.rand(len(prompt)) < self.p_zero).numpy().tolist()
|
||
prompt = [
|
||
self.train_n_prompt if zeros[idx] else p
|
||
for idx, p in enumerate(prompt)
|
||
]
|
||
|
||
with torch.autocast(device_type='cuda', enabled=False):
|
||
context = self.encode_condition(
|
||
self.tokenizer(prompt).to(we.device_id))
|
||
|
||
loss = self.diffusion.loss(x0=x_start,
|
||
t=t,
|
||
model=self.model,
|
||
model_kwargs={'cond': context},
|
||
noise=noise)
|
||
loss = loss.mean()
|
||
ret = {'loss': loss, 'probe_data': {'prompt': prompt}}
|
||
return ret
|
||
```
|
||
|
||
The forward_test function is used during the inference stage to perform the full image denoising process.
|
||
|
||
```python
|
||
@torch.no_grad()
|
||
@torch.autocast('cuda', dtype=torch.float16)
|
||
def forward_test(self,
|
||
prompt=None,
|
||
n_prompt=None,
|
||
sampler='ddim',
|
||
sample_steps=50,
|
||
seed=2023,
|
||
guide_scale=7.5,
|
||
guide_rescale=0.5,
|
||
discretization='trailing',
|
||
run_train_n=True,
|
||
**kwargs):
|
||
g = torch.Generator(device=we.device_id)
|
||
seed = seed if seed >= 0 else random.randint(0, 2**32 - 1)
|
||
g.manual_seed(seed)
|
||
num_samples = len(prompt)
|
||
|
||
n_prompt = default(n_prompt, [self.default_n_prompt] * len(prompt))
|
||
assert isinstance(prompt, list) and \
|
||
isinstance(n_prompt, list) and \
|
||
len(prompt) == len(n_prompt)
|
||
|
||
context = self.encode_condition(self.tokenizer(prompt).to(
|
||
we.device_id), method='encode_text')
|
||
null_context = self.encode_condition(self.tokenizer(n_prompt).to(
|
||
we.device_id), method='encode_text')
|
||
|
||
width, height = 512, 512
|
||
noise = self.noise_sample(num_samples, width // self.size_factor,
|
||
height // self.size_factor, g)
|
||
# UNet use input n_prompt
|
||
samples = self.diffusion.sample(solver=sampler,
|
||
noise=noise,
|
||
model=self.model,
|
||
model_kwargs=[{
|
||
'cond': context
|
||
}, {
|
||
'cond': null_context
|
||
}],
|
||
steps=sample_steps,
|
||
guide_scale=guide_scale,
|
||
guide_rescale=guide_rescale,
|
||
discretization=discretization,
|
||
show_progress=True,
|
||
seed=seed,
|
||
condition_fn=None,
|
||
clamp=None,
|
||
percentile=None,
|
||
t_max=None,
|
||
t_min=None,
|
||
discard_penultimate_step=None,
|
||
return_intermediate=None,
|
||
**kwargs)
|
||
x_samples = self.decode_first_stage(samples).float()
|
||
x_samples = torch.clamp((x_samples + 1.0) / 2.0, min=0.0, max=1.0)
|
||
|
||
outputs = list()
|
||
for p, np, img in zip(prompt, n_prompt, x_samples):
|
||
one_tup = {'prompt': p, 'n_prompt': np, 'image': img}
|
||
outputs.append(one_tup)
|
||
|
||
return outputs
|
||
```
|
||
|
||
# 3. Submodule Registration
|
||
|
||
After the network class is fully implemented, it is necessary to ensure that all submodules used within the class are registered. For instance, with the embedder in SD1.5, to instantiate this embedder within the network's initialization method, we must first implement the embedder class and register it with Scepter.
|
||
|
||
```python
|
||
@EMBEDDERS.register_class()
|
||
class FrozenCLIPEmbedder(BaseEmbedder):
|
||
def __init__(self, cfg, logger=None):
|
||
super().__init__(cfg, logger=logger)
|
||
|
||
with FS.get_dir_to_local_dir(cfg.PRETRAINED_MODEL,
|
||
wait_finish=True) as local_path:
|
||
self.tokenizer = CLIPTokenizer.from_pretrained(local_path)
|
||
self.transformer = CLIPTextModel.from_pretrained(local_path)
|
||
|
||
self.use_grad = cfg.get('USE_GRAD', False)
|
||
self.freeze_flag = cfg.get('FREEZE', True)
|
||
if self.freeze_flag:
|
||
self.freeze()
|
||
|
||
self.max_length = cfg.get('MAX_LENGTH', 77)
|
||
self.layer = cfg.get('LAYER', 'last')
|
||
self.layer_idx = cfg.get('LAYER_IDX', None)
|
||
self.use_final_layer_norm = cfg.get('USE_FINAL_LAYER_NORM', False)
|
||
assert self.layer in self.LAYERS
|
||
if self.layer == 'hidden':
|
||
assert self.layer_idx is not None
|
||
assert 0 <= abs(self.layer_idx) <= 12
|
||
|
||
def freeze(self):
|
||
self.transformer = self.transformer.eval()
|
||
for param in self.parameters():
|
||
param.requires_grad = False
|
||
|
||
# @torch.no_grad()
|
||
def _encode_text(self,
|
||
tokens,
|
||
tokenizer=None,
|
||
append_sentence_embedding=False):
|
||
outputs = self.transformer(input_ids=tokens,
|
||
output_hidden_states=self.layer == 'hidden')
|
||
if self.layer == 'last':
|
||
z = outputs.last_hidden_state
|
||
elif self.layer == 'pooled':
|
||
z = outputs.pooler_output[:, None, :]
|
||
if self.use_final_layer_norm:
|
||
z = self.transformer.text_model.final_layer_norm(z)
|
||
else:
|
||
z = outputs.hidden_states[self.layer_idx]
|
||
if self.use_final_layer_norm:
|
||
z = self.transformer.text_model.final_layer_norm(z)
|
||
return z
|
||
|
||
def encode_text(self,
|
||
tokens,
|
||
tokenizer=None,
|
||
append_sentence_embedding=False):
|
||
if not self.use_grad:
|
||
with torch.no_grad():
|
||
output = self._encode_text(tokens, tokenizer,
|
||
append_sentence_embedding)
|
||
else:
|
||
output = self._encode_text(tokens, tokenizer,
|
||
append_sentence_embedding)
|
||
return output
|
||
|
||
@staticmethod
|
||
def get_config_template():
|
||
return dict_to_yaml('MODELS',
|
||
__class__.__name__,
|
||
FrozenCLIPEmbedder.para_dict,
|
||
set_name=True)
|
||
```
|
||
|
||
# 4. Solver Registration
|
||
|
||
The solver class encapsulates the complete process needed to train and test a network. For instance, to register a solver for training the SD1.5 model, you'd need to:
|
||
|
||
1. Create one or more data loaders, corresponding to the construct_data method in the solver;
|
||
2. Instantiate a model, corresponding to the construct_model method in the solver;
|
||
3. (Optional) Define metrics, corresponding to the construct_metrics method in the solver;
|
||
4. (Optional) Define various hooks (HOOKS) used in training and testing, such as model saving, pre-trained parameter loading, logging, etc.
|
||
|
||
```python
|
||
@SOLVERS.register_class()
|
||
class LatentDiffusionSolver(BaseSolver):
|
||
def set_up(self):
|
||
self.construct_data()
|
||
self.construct_model()
|
||
self.construct_metrics()
|
||
self.model_to_device()
|
||
self.init_opti()
|
||
|
||
def load_checkpoint(self, checkpoint):
|
||
pass
|
||
|
||
def save_checkpoint(self):
|
||
pass
|
||
|
||
def solve(self):
|
||
self.before_solve()
|
||
if 'train' in self._mode_set:
|
||
self.run_train()
|
||
if 'test' in self._mode_set:
|
||
self.run_test()
|
||
self.after_solve()
|
||
|
||
def run_train(self):
|
||
pass
|
||
|
||
def run_eval(self):
|
||
pass
|
||
|
||
def run_test(self):
|
||
pass
|
||
```
|
||
|
||
# 5. Training/Testing Hyperparameter Definitions
|
||
|
||
After registering all required components, which include but are not limited to BACKBONE, NETWORK, EMBEDDER, SOLVER, and METRIC, it is necessary to configure some of the hyperparameters used. Scepter employs YAML files to define hyperparameters for each module; for specifics, refer to scepter/modules/examples/sd15/sd15_512_full.yaml.
|
||
|
||
# 6. Model Training
|
||
|
||
Execute training or batch testing by loading the required YAML file with the --cfg parameter.
|
||
```shell
|
||
# Multi-node Multi-GPU Training
|
||
# Spawn mode (default)
|
||
export CUDA_VISIBLE_DEVICES=0,1,2,3
|
||
export WORLD_SIZE=1
|
||
python -W ignore scepter/run_train.py --cfg scepter/examples/sd15/sd15_512_full.yaml
|
||
# torchrun mode
|
||
torchrun --nproc_per_node 4 scepter/run_train.py --cfg scepter/examples/sd15/sd15_512_full.yaml --launcher torchrun
|
||
# pytorch_lightning mode， ENV.USE_PL=true
|
||
python -W ignore scepter/run_train.py --cfg scepter/examples/sd15/sd15_512_full.yaml
|
||
# Single-GPU Training
|
||
# Spawn mode (default)
|
||
export CUDA_VISIBLE_DEVICES=0
|
||
export WORLD_SIZE=1
|
||
python -W ignore scepter/run_train.py --cfg scepter/examples/sd15/sd15_512_full.yaml
|
||
# torchrun mode
|
||
python -W ignore scepter/run_train.py --cfg scepter/examples/sd15/sd15_512_full.yaml --launcher torchrun
|
||
# pytorch_lightning mode， ENV.USE_PL=true
|
||
python scepter/run_train.py --cfg scepter/examples/sd15/sd15_512_full.yaml
|
||
```
|
||
|
||
# 7. Model Inference
|
||
|
||
Single inference can be realized by customizing the run_inference.py file, with reference to the inference method of SD1.5.
|
||
```shell
|
||
python -W ignore scepter/run_inference.py --prompt "a woman" --n_prompt "" --num_samples 4 --pretrained_model "path/to/your/pretrained/model"
|
||
```
|