Back-End/Python

Python, TensorFlow Diffusion model 학습 및 이미지 생성하기(Image generator), Machine learning, Deep learning for window

개발자 DalBy 2024. 5. 14. 16:52
반응형

Python, TensorFlow Diffusion model 학습 및 이미지 생성하기 (Image generator), Machine learning, Deep learning for window

 

이번 Diffusion model 생성 방법에 대해 포스팅 해보겠습니다. 먼저 Diffusion model은 확산 모델이라고 하며, 확산 과정을 학습하여 새로운 이미지를 생성하는 확산 확률 모델이라고 합니다.

 

이론적인 지식은 관련 논문이나, 관련 유튜브 동영상 또는 아래 위키를 참고하시면 좋습니다.

https://en.wikipedia.org/wiki/Diffusion_model

 

Diffusion model - Wikipedia

From Wikipedia, the free encyclopedia Deep learning algorithm In machine learning, diffusion models, also known as diffusion probabilistic models or score-based generative models, are a class of latent variable generative models. A diffusion model consists

en.wikipedia.org

 

전반적인 확산 모델 DDPM은 GAN, VAE와 같은 다른 생성 모델과 비교 해 보면 그렇게 복잡하지 않다고 합니다.(제가 보기엔 엄청 복잡함) 이들은 모두 단순 분포의 노이즈를 데이터 샘플로 변환합니다. 이것은 신경망이 순수한 노이즈부터 시작하여 점차적으로 데이터의 노이즈를 제거하는 방법을 학습합니다.

 

 

먼저 아래 코드는 Diffusion model의 구현 코드는 아니지만, 간단한 forward 프로세스의 예시를 확인 할 수 있습니다.

import numpy as np
import cv2
import matplotlib.pyplot as plt
import torch

plt.rcParams['font.family'] ='Malgun Gothic'
plt.rcParams['axes.unicode_minus'] =False

# forward process 방식 시각화 테스트
# 디퓨전 모델 테스트 노이즈 레벨
def diffusion_model(image, noise_level):
    noisy_image = image + np.random.normal(scale=noise_level, size=image.shape)
    noisy_image = np.clip(noisy_image, 0, 255).astype(np.uint8)  # [0, 255]

    return noisy_image


# 이미지 로드
image = cv2.cvtColor(cv2.imread("./noiseImg/test.png"), cv2.COLOR_RGB2BGR)

# 노이즈 양
noise_level = 15

# 디퓨전 모델 테스트
plt.figure(figsize=(25, 12))
for i in range(27):
    ax = plt.subplot(3, 9, i+1)
    if i == 0:
        plt.title("원본")
        plt.imshow(image)
    else:
        plt.title("noise level " + str(i))
        noise_level = noise_level + (i * 10)
        noisy_image = diffusion_model(image, noise_level)
        plt.imshow(noisy_image)

    ax.axis("off")

plt.show()

 

저번 시간에 포스팅한 한복을 입은 캐릭터를 이용하여 학습하는 내용을 토대로 테스트를 진행 해 보면 아래 사진과 같은 결과를 확인 할 수 있습니다.

 

테스트 결과:

한복을 입은 캐릭터 테스트
한복을 입은 캐릭터 테스트

테스트 결과를 보면 점점 미지 파일에 노이즈가 점차 생기는 것을 확인 할 수 있습니다. 이러한 방식을 이용하여 DDPM을 구현 해 보려고 합니다. 전반적인 확산 모델의 기본 지식을 참고 하였다면 실습을 통해 구현하는 방법을 알아보겠습니다. (단 참고사항으로 DDPM의 수학 공식을 구현해야 합니다.)

 

 

먼저 코드를 작성하기 전 pytorch를 import 해야 합니다.

https://pytorch.kr/get-started/locally/

 

파이토치 한국 사용자 모임 (PyTorch Korea User Group)

파이토치 한국 사용자 모임에 오신 것을 환영합니다. 딥러닝 프레임워크인 파이토치(PyTorch)를 사용하는 한국어 사용자들을 위해 문서를 번역하고 정보를 공유하고 있습니다.

pytorch.kr

 

필자가 설치한 버전은 아래와 같습니다. ( 필자의 개발 환경 사항은 이전 가이드 포스팅을 참고해 주세요 ! )

pip install torch==1.7.1+cu110 torchvision==0.8.2+cu110 -f https://download.pytorch.org/whl/torch_stable.html

 

pytorch 설치 후, 정상적으로 설치 되었는지 확인합니다.

import torch
torch.cuda.is_available()

=true

 

또는

x = torch.rand(5, 3)
print(x)

출력
tensor([[0.6908, 0.6596, 0.6259], 
        [0.4186, 0.9298, 0.7164], 
        [0.3510, 0.2787, 0.7555], 
        [0.2269, 0.9852, 0.1033], 
        [0.7045, 0.5437, 0.1077]])

 

 

먼저 U-net architecture를 이해하고, 구현해야 합니다.

U-net architecture 이미지
U-net architecture 이미지

설명은 현 포스팅에서 생략하고 관련 링크는 다음과 같습니다.

https://paperswithcode.com/method/u-net

 

Papers with Code - U-Net Explained

U-Net is an architecture for semantic segmentation. It consists of a contracting path and an expansive path. The contracting path follows the typical architecture of a convolutional network. It consists of the repeated application of two 3x3 convolutions (

paperswithcode.com

 

먼저 module.py 파일을 생성하여 코드를 구현하면 다음과 같습니다.

import torch
import torch.nn as nn
import torch.nn.functional as F

# Exponential Moving Average
# 지수 이동 평균 알고리즘
class EMA:
    def __init__(self, beta):
        super().__init__()
        self.beta = beta
        self.step = 0

'''
    MultiheadAttention을 사용하여 Attention 가중치를 계산
    다음 레이어 정규화를 적용하고 마지막 피드 포워드 네트워크를 통해 결과 전달
    (해당 모델이 입력 시퀀스 관련 부분의 표현을 개선하는것에 도움이 된다고 한다.)
'''
class SelfAttention(nn.Module):
    # (self, 입력 채널, attention size)
    def __init__(self, chs, size):
        super(SelfAttention, self).__init__()
        self.chs = chs
        self.size = size

         # Multihead로 분할하고 각 head에 대해 독립적으로 스케일링된 내적 주의를 계산하는 모듈을 생성 
        self.mha = nn.MultiheadAttention(chs, 4)

        # MultiheadAttention 계산 모듈 적용
        self.multiheadAttention = nn.MultiheadAttention(chs, 4)

        # 레이어 정규화 MultiheadAttention에 입력 받기 전 입력 레이어 정규화 적용
        # 이것을 적용하면 학습의 안정화에 대한 도움이 될 수도 있다고 한다. 
        self.layerNorm = nn.LayerNorm([chs])

         # Multihead 피드 포워드 신경망 정의  
        self.f_seq = nn.Sequential(
            nn.LayerNorm([chs]), # 입력 레이어 정규화
            nn.Linear(chs, chs), # 선형 레이어
            nn.GELU(), 
            nn.Linear(chs, chs)
        )

    # 정방향 프로세스 
    def forward(self, x):
        x = x.view(-1, self.chs, self.size * self.size).transpose(1, 2)
        x_layerNorm = self.layerNorm(x) # 레이어 정규화를 통해 전달
        att_val, _ = self.mha(x_layerNorm, x_layerNorm, x_layerNorm)
        att_val = att_val + x
        att_val = self.f_seq(att_val) + att_val # 신경망 통과 비선형성이 도입 -> 표현을 구체화 
        return att_val.transpose(2, 1).view(-1, self.chs, self.size, self.size)

 

SelfAttention class는 MultiheadAttention을 사용하여 Attention가충치를 계산, 다음 레이어 정규화를 적용하고 마지막 피드 포워드 네트워크를 통해 결과를 전달 합니다. 해당 모델이 입력 시퀀스 관련 부분의 표현을 개선하는 것에 도움이 된다고 합니다. __init__ 메소드와 정방향 프로세스 forward 메소드를 구현합니다. 

 

'''
    이중 컨볼루션 레이어와 정규화 및 활성화 함수로 구성된 모듈을 정의 U-Net 아키텍처 빌딩 블록
'''
class DoubleConv(nn.Module):
    def __init__(self, in_chs, out_chs, mid_chs=None, residual=False):
        super().__init__()
        self.residual = residual
        if not mid_chs:
            mid_chs = out_chs

         # 입력 채널, 출력 채널, 3x3 커널 크기 및 1픽셀 패딩이 포함된 첫 번째 컨벌루션 레이어
         # 채널 차원을 따라 적용
        self.double_conv = nn.Sequential(
            nn.Conv2d(in_chs, mid_chs, kernel_size=3, padding=1, bias=False),
            nn.GroupNorm(1, mid_chs),
            nn.GELU(),
            nn.Conv2d(mid_chs, out_chs, kernel_size=3, padding=1, bias=False),
            nn.GroupNorm(1, out_chs)
        )

    # 정방향 프로세스 
    def forward(self, x):
         
        if self.residual:
            return F.gelu(x + self.double_conv(x))
        else:
            return self.double_conv(x)

그 다음 위 코드의 이중 컨볼루션 레이어와 정규화 및 활성화 함수로 구성된 모듈을 정의 합니다.

 

'''
    다운 샘플링 구성 (공간 해상도를 줄이는 신경망의 다운 샘플링)
'''
class Down(nn.Module):
    def __init__(self, in_chs, out_chs, embedding_dim=256):
        super().__init__()
        self.down_maxpool_conv = nn.Sequential(
            nn.MaxPool2d(2), # 맥스풀링 
            DoubleConv(in_chs, in_chs, residual=True),
            DoubleConv(in_chs, out_chs)
        )

        # 임베딩 레이어 
        self.embedding_layer = nn.Sequential(
            nn.SiLU(),
            nn.Linear(
                embedding_dim,
                out_chs
            ),
        )

    # 정방향 프로세스
    def forward(self, x, t):
        x = self.down_maxpool_conv(x)
        embedding = self.embedding_layer(t)[:, :, None, None].repeat(1, 1, x.shape[-2], x.shape[-1])
        return x + embedding

위 코드를 참고하여 공간 해상도를 줄이는 다운 샘플링을 구현합니다.

 

'''
    업 샘플링 구성 (공간 해상도를 높이는 신경망의 업샘플링)
'''
class Up(nn.Module):
    def __init__(self, in_chs, out_chs, embedding_dim=256):
        super().__init__()

        self.up = nn.Upsample(scale_factor=2, mode="bilinear", align_corners=True) # 쌍선형 보간법을 사용하여 입력 텐서를 2배로 업샘플링하는 모듈 
        self.up_conv = nn.Sequential(
            DoubleConv(in_chs, in_chs, residual=True),
            DoubleConv(in_chs, out_chs, in_chs // 2)
        )

        # 임베딩 레이어
        self.embedding_layer = nn.Sequential(
            nn.SiLU(),
            # 선형 변환 
            nn.Linear(
                embedding_dim,
                out_chs
            ),
        )

    # 정방향 프로세스
    def forward(self, x, skip_x, t):
        x = self.up(x)
        x = torch.cat([skip_x, x], dim=1)
        x = self.up_conv(x)
        embedding = self.embedding_layer(t)[:, :, None, None].repeat(1, 1, x.shape[-2], x.shape[-1])
        return x + embedding

그 다음 위 코드 참고하여 공간 해상도를 높이는 업샘플링을 구현합니다.

 

 

'''
    unet 구현 : RGB 이미지로 작업하고 있으므로 입력 및 출력 채널이 둘다 3이 될 유닛에 대한 기본 구성을 정의한다
    c_in 입력 채널 수 RGB 3
    c_out 출력 채널 수 3
    time_dim 시간 차원의 기본 차원값 256
    device 모델 실행 장치 GPU
'''
class UNet(nn.Module):
    def __init__(self, in_ch=3, out_ch=3, time_dim=256, device="cuda"):
        super().__init__()
        self.device = device
        self.time_dim = time_dim

        '''
            인코더 Down() 블록과 인터리브 된 다운샘플링 블록으로 구성
            inc 초기 컨볼루션 수행
            down1 ~ 3 채널 수를 늘리면서 공간 차원을 점진적으로 다운 샘플링
            전역 종속성을 캡처하기 위해 사용 sa1 ~ 3
        '''
        self.inc = DoubleConv(in_ch, 64)
        self.down1 = Down(64, 128)
        self.sa1 = SelfAttention(128, 32)
        self.down2 = Down(128, 256)
        self.sa2 = SelfAttention(256, 16)
        self.down3 = Down(256, 256)
        self.sa3 = SelfAttention(256, 8)

        # 타임스탬프 정보 포함, 임베딩 차원 (명시적 포함 x)
        self.bot1 = DoubleConv(256, 512)
        self.bot2 = DoubleConv(512, 512)
        self.bot3 = DoubleConv(512, 256)

        '''
            전반적으로 이 U-Net 아키텍처는 표준 인코더-디코더 구조 외에도 장거리 종속성을 캡처하기 위한 self-attention 메커니즘을 통합
            채널 수를 줄이면서 업샘플링을 수행 sa4~6
            디코더 Up() 업샘플링 블록
            출력 레이어 outc는 최종 컨볼루션을 수행하여 c_out 채널이 있는 분할 마스크를 생성
        '''
        self.up1 = Up(512, 128)
        self.sa4 = SelfAttention(128, 16)
        self.up2 = Up(256, 64)
        self.sa5 = SelfAttention(64, 32)
        self.up3 = Up(128, 64)
        self.sa6 = SelfAttention(64, 64)
        self.outc = nn.Conv2d(64, out_ch, kernel_size=1)


    '''
        타임스탬프에 대한 위치 인코딩 생성 
        (위치 인코딩은 자연어 처리 또는 시계열 분석과 같은 시퀀스 모델링 작업에서 시퀀스 내의 토큰 또는 타임 스탬프 위치에 대한 정보를 모델에 제공)
        위치 인코딩이 시퀀스에서 타임스탬프의 순서와 상대 위치 정보를 모두 캡처할 수 있으며, 모델이 데이터의 순차적 특성을 이해하기 위한 힌트 같은거인듯?
    '''
    def pos_encoding(self, t, chs):
        # 역 주파수 계산 1 / (10000(i / 채널))
        inv_freq = 1.0 / (
            10000
            ** (torch.arange(0, chs, 2, device=self.device).float() / chs)
        )

        # 사인, 코사인 인코딩
        pos_enc_a = torch.sin(t.repeat(1, chs // 2) * inv_freq)
        pos_enc_b = torch.cos(t.repeat(1, chs // 2) * inv_freq)

        # 채널 차원()을 따라 코사인 인코딩을 연결, dim=-1 각 타임 스탬프 최종 위치에서 인코딩 생성
        pos_enc = torch.cat([pos_enc_a, pos_enc_b], dim=-1)
        return pos_enc
    


    # 정방향 프로세스
    def forward(self, x, t):
        # 전처리 타임스탬프 마지막 차원을 따라 해당 유형으로 변환
        t = t.unsqueeze(-1).type(torch.float)
        t = self.pos_encoding(t, self.time_dim)

        # 인코더 다운 샘플링
        x1 = self.inc(x)
        x2 = self.down1(x1, t)
        x2 = self.sa1(x2)
        x3 = self.down2(x2, t)
        x3 = self.sa2(x3)
        x4 = self.down3(x3, t)
        x4 = self.sa3(x4)

        # 컨볼루션 레이어
        x4 = self.bot1(x4)
        x4 = self.bot2(x4)
        x4 = self.bot3(x4)

        # 디코더 업 샘플링
        x = self.up1(x4, x3, t)
        x = self.sa4(x)
        x = self.up2(x, x2, t)
        x = self.sa5(x)
        x = self.up3(x, x1, t)
        x = self.sa6(x)

        # 특징 전달
        output = self.outc(x)
        return output

그리고 위 코드를 참고하여  U-Net을 구현합니다. 

 

'''
이 U-Net_conditional클래스는 조건부 정보, 특히 추가 입력을 통해 기존 U-Net 아키텍처를 확장합니다 y. 작동 방식은 다음과 같습니다
이 아키텍처를 사용하면 모델이 로 표시되는 정보를 조건부로 통합할 수 있으며 y, 
이는 조건부 이미지 생성이나 추가 레이블이 제공되는 분할 작업과 같이 조건부 정보를 사용할 수 있는 다양한 작업에 도움이 될 수 있습니다.
'''
class UNet_conditional(nn.Module):
    def __init__(self, in_ch=3, out_ch=3, time_dim=256, num_classes=None, device="cuda"):
        super().__init__()
        self.device = device
        self.time_dim = time_dim

        self.inc = DoubleConv(in_ch, 64)
        self.down1 = Down(64, 128)
        self.sa1 = SelfAttention(128, 32)
        self.down2 = Down(128, 256)
        self.sa2 = SelfAttention(256, 16)
        self.down3 = Down(256, 256)
        self.sa3 = SelfAttention(256, 8)

        self.bot1 = DoubleConv(256, 512)
        self.bot2 = DoubleConv(512, 512)
        self.bot3 = DoubleConv(512, 256)

        self.up1 = Up(512, 128)
        self.sa4 = SelfAttention(128, 16)
        self.up2 = Up(256, 64)
        self.sa5 = SelfAttention(64, 32)
        self.up3 = Up(128, 64)
        self.sa6 = SelfAttention(64, 64)
        self.outc = nn.Conv2d(64, out_ch, kernel_size=1)

        if num_classes is not None:
            self.label_emb = nn.Embedding(num_classes, time_dim)

    # 위치 인코딩 
    def pos_encoding(self, t, channels):

        inv_freq = 1.0 / (
            10000
            ** (torch.arange(0, channels, 2, device=self.device).float() / channels)
        )

        pos_enc_a = torch.sin(t.repeat(1, channels // 2) * inv_freq)
        pos_enc_b = torch.cos(t.repeat(1, channels // 2) * inv_freq)
        pos_enc = torch.cat([pos_enc_a, pos_enc_b], dim=-1)
        return pos_enc

    # 정방향 프로세스
    def forward(self, x, t, y):
        t = t.unsqueeze(-1).type(torch.float)
        t = self.pos_encoding(t, self.time_dim)

        if y is not None:
            t += self.label_emb(y)

        x1 = self.inc(x)
        x2 = self.down1(x1, t)
        x2 = self.sa1(x2)
        x3 = self.down2(x2, t)
        x3 = self.sa2(x3)
        x4 = self.down3(x3, t)
        x4 = self.sa3(x4)

        x4 = self.bot1(x4)
        x4 = self.bot2(x4)
        x4 = self.bot3(x4)

        x = self.up1(x4, x3, t)
        x = self.sa4(x)
        x = self.up2(x, x2, t)
        x = self.sa5(x)
        x = self.up3(x, x1, t)
        x = self.sa6(x)
        output = self.outc(x)
        return output

마지막으로 U-Net_conditional 클래스를 구현합니다.

 

이제 주요 기능을 만들었으니, 주요 기능과 Model 학습 코드를 작성할 ddpm.py 파일을 생성하여 코드를 작성하겠습니다.

class Diffusion:
    
    # 이미지 저장
    def save_img(self, images, path, **kwargs):
        grid = torchvision.utils.make_grid(images, **kwargs)
        ndarr = grid.permute(1, 2, 0).to('cpu').numpy()
        im = Image.fromarray(ndarr)
        im.save(path)

    # 데이터 로더
    def get_data(self, params):
        transforms = torchvision.transforms.Compose([
            torchvision.transforms.Resize(80),
            torchvision.transforms.RandomResizedCrop(params.image_size, scale=(0.8, 1.0)),
            torchvision.transforms.ToTensor(),
            torchvision.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
        ])
        dataset = torchvision.datasets.ImageFolder(params.dataset_path, transform=transforms)
        dataloader = DataLoader(dataset, batch_size=params.batch_size, shuffle=True)
        return dataloader
    
    # 실행
    def __init__(self, noise_steps=1000, beta_start=1e-4, beta_end=0.02, img_size=256, device="cuda"):

        self.noise_steps = noise_steps
        self.beta_start = beta_start
        self.beta_end = beta_end
        self.img_size = img_size
        self.device = device

        self.beta = self.noise_schedule().to(device)
        self.alpha = 1. - self.beta
        self.alpha_c_p = torch.cumprod(self.alpha, dim=0)


    # 노이즈 스케쥴러
    def noise_schedule(self):
        return torch.linspace(self.beta_start, self.beta_end, self.noise_steps)
    

    # 노이즈 이미지, 매개변수 기반 이미지에 노이즈를 주입
    def noise_images(self, x, t):
        # self.alpha_c_p[t]의 제곱근을 계산 1 - self.alpha_c_p[t]
        sqrt_alpha = torch.sqrt(self.alpha_c_p[t])[:, None, None, None]
        sqrt_one_minus_alpha = torch.sqrt(1 - self.alpha_c_p[t])[:, None, None, None]
        # e 동일한 형태의 랜덤 노이즈를 생성
        e = torch.randn_like(x) 
        return sqrt_alpha * x + sqrt_one_minus_alpha * e, e
    
    # 시간 단계를 샘플링 (시간적 노이즈 주입, 임의의 시간 단계를 생성)
    def sample_timesteps(self, n):
        return torch.randint(low=1, high=self.noise_steps, size=(n,))

    def model_sample(self, model, n):
        # 드롭아웃 및 배치 정규화 레이어를 비활성화 하여 모델을 평가 모드로 설정
        model.eval()

        with torch.no_grad():
            #  임의의 입력 이미지를 생성합니다 
            x = torch.randn((n, 3, self.img_size, self.img_size)).to(self.device)

            '''
                self.noise_steps - 11(0 제외) 의 역방향 범위를 반복
                각 시간 단계에 대해 의 값을 포함하는 i텐서를 생성합 .t i
                예측된 노이즈를 얻기 위해 입력 이미지 x와 시간 단계 텐서를 사용하여 모델을 호출
            '''
            for i in tqdm(reversed(range(1, self.noise_steps)), position=0):
                t = (torch.ones(n) * i).long().to(self.device)
                predicted_noise = model(x, t)
                alpha = self.alpha[t][:, None, None, None]
                alpha_c_p = self.alpha_c_p[t][:, None, None, None]
                beta = self.beta[t][:, None, None, None]
                if i > 1:
                    noise = torch.randn_like(x)
                else:
                    noise = torch.zeros_like(x)

                # alpha, alpha_c_p, beta, 예측 노이즈 및 랜덤 노이즈를 기반으로 입력 이미지를 조작
                x = 1 / torch.sqrt(alpha) * (x - ((1 - alpha) / (torch.sqrt(1 - alpha_c_p))) * predicted_noise) + torch.sqrt(beta) * noise


        # 드롭아웃 및 배치 정규화 레이어를 활성화 하여 모델을 훈련 모드로 다시 설정
        model.train()

        # 픽셀 값을 [-1, 1] 범위로 고정하고 이를 uint8 형식으로 변환하는 등의 사후 처리 단계를 수행
        x = (x.clamp(-1, 1) + 1) / 2
        x = (x * 255).type(torch.uint8)
        return x

위 코드와 같이 실행 함수와 이미지 저장, 데이터 로더, 노이즈 스케쥴러, 타임스탭, 모델 생성 메소드를 작성 해 줍니다.

 

def train(params):

    diff = Diffusion()

    device = params.device
    dataloader = diff.get_data(params)
    model = UNet().to(device)
    optimizer = optim.AdamW(model.parameters(), lr=params.lr)
    mse = nn.MSELoss()
    diffusion = Diffusion(img_size=params.image_size, device=device)
    l = len(dataloader)

    for epoch in range(params.epochs):
        print(f"epoch {epoch}:")
        pbar = tqdm(dataloader)
        for i, (images, _) in enumerate(pbar):
            images = images.to(device)

            # 배치의 각 이미지에 대한 확산 모델의 시간 단계를 샘플링
            t = diffusion.sample_timesteps(images.shape[0]).to(device)

            # 확산 모델을 사용하여 이미지에 노이즈를 추가
            x_t, noise = diffusion.noise_images(images, t)

            # 잡음이 있는 이미지에 대해 U-Net 모델로부터 예측된 잡음을 얻습니다
            predicted_noise = model(x_t, t)
            
            # 예측된 잡음과 실제 잡음 사이의 MSE 손실을 계산
            loss = mse(noise, predicted_noise)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            # 모델 매개변수를 업데이트하기 위해 역전파 및 최적화를 수행
            pbar.set_postfix(MSE=loss.item())

        # 각 에포크 후에 확산 모델과 훈련된 U-Net 모델을 사용하여 이미지를 샘플링
        sampled_images = diffusion.model_sample(model, n=images.shape[0])
        # 샘플링된 이미지를 디스크에 저장
        diff.save_img(sampled_images, os.path.join("result", params.run_name, f"{epoch}.jpg"))
        
        # 모델 체크포인트 저장
        torch.save(model.state_dict(), os.path.join("model", params.run_name, f"ckpt.pt"))

에포크 진행 후, 확산을 학습 한 모델을 저장하는 메소드를 구현 해 줍니다.

 

실행문은 다음과 같습니다.

def launch():
    import argparse
    parser = argparse.ArgumentParser()
    params = parser.parse_args()
    params.run_name = "DDPM"
    params.epochs = 300
    params.batch_size = 10
    params.image_size = 64  #64
    params.dataset_path = r"C:\\Users\\사용자이름\\.keras\\datasets\\education_images"
    params.device = "cuda"
    params.lr = 3e-4
    train(params)

if __name__ == '__main__':
    launch()

run_name의 경우 테스트 내용이 저장될 기준 폴더가 됩니다. 학습하고 싶은 내용에 따라 batch_size, epochs 값을 입력 해주고 학습 할 이미지는 64*64 크기의 이미지 파일로 진행하였습니다.

 

(참고사항)

처음 batch size는12정도, 에포크(epochs)는 500정도 테스트 진행 해 보았습니다. RTX3060TI기준 대략 7시간 정도가 소요되었습니다.

 

테스트 결과는 다음과 같습니다. 

총 이미지 파일 700개 이상 테스트를 진행하였고, batch_size = 12,  epochs = 500로 하여 진행하였습니다. 체크포인트 파일용량은 89MB 정도로 생성되었습니다.

 

다음 한복을 입은 캐릭터 학습 테스트 결과에 대한 참고 이미지 입니다.(테스트 데이터는 novel AI의 이미지 제너레이터를 이용하여 생성하였습니다. 해당 AI 버전은 NAI V3 버전입니다.)

 

1. epochs 0, 30, 50, 70번째 생성된 이미지 파일 입니다.

epochs 0~70번째
epochs 0~70번째

해당 모델이 이미지를 생성 한 후, 이미지 파일을 확인 해 보면 노이즈만 가득한 것을 알 수 있습니다. 아메바 처럼 생겼습니다...

 

 

 

2. epochs 90, 110, 130, 150번째 생성된 이미지 파일 입니다.

epochs 90~150번째
epochs 90~150번째

이제 조금씩 캐릭터 형태의 윤곽이 잡히는 것 같습니다. 뭔가 검고, 빨간 이미지 중, 기괴하고 소름 돋았던 이미지 파일도 있었습니다.

 

 

 

3. epochs 170, 190, 210, 230번째 생성된 이미지 파일 입니다.

epochs 170~230번째
epochs 170~230번째

이때부터는 캐릭터의 형태가 뚜렷하지만, 머리가 두개 있거나, 샴쌍둥이(?) 같은 기괴한(?) 캐릭터가 많이 생성되었습니다. 느낌이 많이 무섭습니다. ㄷㄷ

 

 

4. epochs 410, 430, 450, 497번째 생성된 이미지 파일 입니다.

 

epochs 410~497번째
epochs 410~497번째

학습을 거의 마치기 전 단계까지 왔을 때, 캐릭터 같은(?) 이미지 파일로 생성되었습니다.

 

 

5. 마지막으로 생성된 한복을 입은 캐릭터 이미지 데이터 입니다.

epochs 499
epochs 499

해당 모델의 체크포인트를 지속적으로 학습한다면 여러 다양한 방면으로 AI를 활용할 수 있을 것 같습니다.

 

 



반응형