down_unet.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import math
import sympy as sp
import wandb
from PIL import Image
from datasets import load_dataset
from torchvision import transforms



device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device)

# 初始化项目
wandb.init(
    # set the wandb project where this run will be logged
    project="unet-try",
)

'''

conv_block = resnetblock--attentionblock--convblock.   input:[B,C,H,W],output:[B,channel_dim,H(+/-)2,W(+/-)2]
                    
down block = 2blocks|-->for_skip_connection
                    |
                    down_sample-->result_after_pool.    input:[B,C,H,W],output:[B,channel_dim,(H-4)//2,(W-4)//2]
                    
up block = -->concat-->2blocks         input:[B,C,H,W],input_skip:[B,C/2,2H,2W],output:[B,C/2,2H+4,2W+4]
              |
     --up_sample

LR-----------------------------MSE LOSS--------------------------LR
|--down block -------------skip connection-----------up block--|
             |--down block                 up block--|
                          |---------------|
'''

# ----------------------------------------------------------------------------------------------------
class conv_block(nn.Module):    #一个下采样模块包含两个卷积层，深度channel从1-64-128-256这样[B,C,H,W]-->[B,C_DIM,H-2,W-2]
    def __init__(self,in_channel,num_heads,channel_dim,use ="down"):
        super(conv_block,self).__init__()      #in_channel输入通道数，channle_dim输出通道数,一个块减少2


        self.in_channel = in_channel
        self.num_heads = num_heads
        self.channel_dim = channel_dim
        self.use = use

        self.GN = nn.GroupNorm(num_groups=4, num_channels=in_channel)  #这个channel指的是输入通道数
        # num_groups 是组数(2,4,8)输入特征的通道分成多少组进行归一化，num_channels 是输入的通道数
        self.conv = nn.Conv2d(in_channels=in_channel, out_channels=in_channel, kernel_size=3,
                          stride=1, padding=1, bias=False)
        self.silu = nn.SiLU()
        self.attention = nn.MultiheadAttention(embed_dim=self.in_channel, num_heads=self.num_heads)

        if self.use == "down":
            self.conv1 = nn.Conv2d(in_channels=self.in_channel, out_channels=self.channel_dim, kernel_size=3,
                             stride=1, padding=0, bias=False)
        elif self.use =="up":
            self.conv1 = nn.Conv2d(in_channels=self.in_channel, out_channels=self.channel_dim, kernel_size=3,
                         stride=1, padding=2, bias=False)



    def resnet_block(self,X):    #隐藏层使用和输入一样的大小

        out = self.GN(X)
        out = self.conv(out)
        out = self.silu(out)

        out = self.GN(out)
        out = self.conv(out)
        out = self.silu(out)

        return out + X

    def attention_block(self,X):

        B,C,H,W = X.size()

        out = self.GN(X)
        out = self.conv(out)

        out = out.view(B, self.in_channel, H * W).transpose(1, 2)  # 将输入重构为 [B, L, C]，其中 L = H * W
        out, weights = self.attention(out, out, out)
        out = out.transpose(1, 2).view(B, self.in_channel, H, W)

        out = self.conv(out)

        return out+X

    def forward(self,X):

        out = self.resnet_block(X)
        out = self.attention_block(out)
        out = self.conv1(out)

        return out


'''
model = conv_block(in_channel=4,num_heads=4,channel_dim=64,use="down")
in_put = torch.randn(1,4,256,256)     #注意，在SR3代码中隐藏层是不变的和输入一致
ouput = model(in_put)
print(ouput.shape)
'''
# -------------------------------------------------------------------------------------------------
class SpatialAttention(nn.Module):
    def __init__(self, in_channels):
        super(SpatialAttention, self).__init__()
        self.conv = nn.Conv2d(in_channels, 1, kernel_size=1)

    def forward(self, x):
        # Apply convolution to generate attention map
        attention_map = self.conv(x)
        # Generate attention scores
        attention_scores = torch.softmax(attention_map, dim=1)
        # Apply attention scores
        out = x * attention_scores
        return out

class ChannelAttention(nn.Module):
    def __init__(self, in_channels, reduction_ratio=16):
        super(ChannelAttention, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Sequential(
            nn.Linear(in_channels, in_channels // reduction_ratio, bias=False),
            nn.ReLU(),
            nn.Linear(in_channels // reduction_ratio, in_channels, bias=False),
            nn.ReLU()
        )

    def forward(self, x):
        # Average pooling to generate a channel descriptor
        avg_out = self.avg_pool(x).view(x.size(0), -1)
        # Apply fully connected layers to generate channel attention
        attn = self.fc(avg_out)
        # Reshape attention to match the input
        attn = attn.view(x.size(0), -1, 1, 1)
        return x * attn


def calculate_attention(X, num_heads, use):
    X = X.to(device)
    B, C, H, W = X.size()

    if use == "down":
        # Apply channel attention
        channel_attention = ChannelAttention(C).to(device)
        out = channel_attention(X)
    elif use == "up":
        # Reshape and transpose for multi-head attention
        up = X.view(B, C, H * W).transpose(1, 2)
        spatial_attention = nn.MultiheadAttention(embed_dim=C, num_heads=num_heads).to(device)
        out, weights = spatial_attention(up, up, up)
        # Apply spatial attention on upsampled output
        out = out.transpose(1, 2).view(B, C, H,W)
        spatial_attention_module = SpatialAttention(in_channels=C).to(device)
        out = spatial_attention_module(out)
        # Reshape output to match the original input dimensions


    return out

'''
# Example usage
X = torch.randn(1,4,572,572)  # Example input tensor
num_heads = 4
attention_out = calculate_attention(X, num_heads,use="up")
print("attention out",attention_out.shape)
'''
'''
X = torch.randn(1, 64, 254, 254)
output = calculate_attention(X,num_heads=8)
print("attention", output.shape)  # 应该输出 torch.Size([1, 64, 254, 254])
'''
# -----------------------------------------------------------------------------------

def generate_positional_encoding(X):
    X = X.to(device)
    B,C,H,W = X.size()
    # 初始化位置编码矩阵
    pos_encoding = torch.zeros(B, C, H, W)

    # 计算位置索引
    y_positions = torch.arange(0, H, dtype=torch.float32).unsqueeze(1).repeat(1, W) #[H,W]
    x_positions = torch.arange(0, W, dtype=torch.float32).unsqueeze(0).repeat(H, 1)

    # 将位置索引除以尺度以进行缩放
    y_positions = y_positions / (H ** 0.5)
    x_positions = x_positions / (W ** 0.5)

    # 计算位置编码的正弦和余弦值
    for i in range(0, C, 2):
        pos_encoding[:, i, :, :] = torch.sin(x_positions)
        pos_encoding[:, i + 1, :, :] = torch.cos(y_positions)

    return pos_encoding

'''
X = torch.randn(1,128, 512, 512)
# 计算位置编码
pos_encoding = generate_positional_encoding(X)
print("Positional Encoding shape:", pos_encoding.shape)  # 应该输出 torch.Size([1, 64, 254, 254])
'''

class down_block(nn.Module):  #宽高减4，然后除以2
    def __init__(self,in_channel,channel_dim):
        super(down_block,self).__init__()

        self.channel_dim = channel_dim

        self.block1 = conv_block(in_channel=in_channel,num_heads=4,
                           channel_dim=self.channel_dim,use="down")
        self.block2 = conv_block(in_channel=self.channel_dim, num_heads=4,
                                channel_dim=self.channel_dim, use="down")

        self.down_pool = nn.Conv2d(in_channels=self.channel_dim, out_channels=self.channel_dim, kernel_size=2,
                              stride=2, padding=0, bias=False)


    def forward(self,X):  #输入[1,4,128,128],输出[1.64,124,124]-->[1,64,62,62]

        out = self.block1(X)
        for_skip_connection = self.block2(out)  #这个out用于跳跃连接的
        result_after_pool = self.down_pool(for_skip_connection)

        return result_after_pool,for_skip_connection

'''
model1 = down_block(in_channel=64,channel_dim=128)
input = torch.randn(1,64,284,284)
res,out = model1(input)
print(res.shape,out.shape)
'''
# --------------------------------------------------------------------------------------------------
class up_block(nn.Module):
    def __init__(self,in_channel):
        super(up_block,self).__init__()
        self.in_channel = in_channel


        self.block1 = conv_block(in_channel=in_channel*2, num_heads=4,
                           channel_dim=in_channel,use="up")
        self.block2 = conv_block(in_channel=in_channel, num_heads=4,
                           channel_dim=in_channel,use="up")
        self.up_pool = nn.ConvTranspose2d(self.in_channel*2, self.in_channel,
                                          kernel_size=2, stride=2)

    def forward(self,input,input_skip):  #先对输入进行上采样，然后和跳跃的拼接，之后经过两个block


        after_transposed = self.up_pool(input)   #上采样得到的大小
        after_cat = torch.cat((after_transposed, input_skip), dim=1)  # 拼接张量
        out = self.block1(after_cat)
        out = self.block2(out)

        return out,after_transposed

'''
model2 = up_block(in_channel=128)
input = torch.randn(1,256,140,140)
input_skip = torch.randn(1,128,280,280)
out,after = model2(input,input_skip)
print("up block",out.shape)  #torch.Size([1, 128, 284, 284])
'''


class down_model(nn.Module):
    def __init__(self):
        super(down_model,self).__init__()

        self.start_conv = nn.Conv2d(in_channels=3, out_channels=4, kernel_size=1, stride=1)

        self.down_block1 = down_block(4,64)
        self.down_block2 = down_block(64,128)
        self.down_block3 = down_block(128,256)
        self.down_block4 = down_block(256,512)

        self.bottle_conv = nn.Conv2d(in_channels=512, out_channels=1024, kernel_size=1, stride=1)

        self.up_block4 = up_block(512)
        self.up_block3 = up_block(256)
        self.up_block2 = up_block(128)
        self.up_block1 = up_block(64)

        self.final_conv = nn.Conv2d(in_channels=64, out_channels=3, kernel_size=1, stride=1)

    def forward(self,input):   #这个地方的输入一定要除的尽

        input = self.start_conv(input)

        result_after_pool1, for_skip_connection1 = self.down_block1(input)
        attention_out1 = calculate_attention(for_skip_connection1, num_heads=4, use="down")
        pos_encoding1 = generate_positional_encoding(for_skip_connection1)
        # print("1",result_after_pool1.shape,for_skip_connection1.shape)


        result_after_pool2, for_skip_connection2 = self.down_block2(result_after_pool1)
        attention_out2 = calculate_attention(for_skip_connection2, num_heads=4, use="down")
        pos_encoding2 = generate_positional_encoding(for_skip_connection2)
        # print("2",result_after_pool2.shape, for_skip_connection2.shape)

        result_after_pool3, for_skip_connection3 = self.down_block3(result_after_pool2)
        attention_out3 = calculate_attention(for_skip_connection3, num_heads=4, use="down")
        pos_encoding3 = generate_positional_encoding(for_skip_connection3)
        # print("3",result_after_pool3.shape, for_skip_connection3.shape)

        result_after_pool4, for_skip_connection4 = self.down_block4(result_after_pool3)
        attention_out4 = calculate_attention(for_skip_connection4, num_heads=4, use="down")
        pos_encoding4 = generate_positional_encoding(for_skip_connection4)
        # print("4",result_after_pool4.shape, for_skip_connection4.shape)


        result_after_pool4 = self.bottle_conv(result_after_pool4)
        # print("bottle",result_after_pool4.shape)


        out, after_transposed1 = self.up_block4(result_after_pool4, for_skip_connection4)
        attention_out5 = calculate_attention(after_transposed1, num_heads=4, use="up")
        pos_encoding5 = generate_positional_encoding(after_transposed1)
        # print("5",out.shape,after_transposed1.shape)


        out, after_transposed2 = self.up_block3(out, for_skip_connection3)
        attention_out6 = calculate_attention(after_transposed2, num_heads=4, use="up").to(device)
        pos_encoding6 = generate_positional_encoding(after_transposed2).to(device)
        # print("6",out.shape, after_transposed2.shape)


        out, after_transposed3 = self.up_block2(out, for_skip_connection2)
        attention_out7 = calculate_attention(after_transposed3, num_heads=4, use="up").to(device)
        pos_encoding7 = generate_positional_encoding(after_transposed3).to(device)
        # print("7",out.shape, after_transposed3.shape)


        out, after_transposed4 = self.up_block1(out, for_skip_connection1)
        attention_out8 = calculate_attention(after_transposed4, num_heads=4, use="up").to(device)
        pos_encoding8 = generate_positional_encoding(after_transposed4).to(device)
        # print("8",out.shape, after_transposed4.shape)


        out = self.final_conv(out)

        return out,attention_out1,attention_out2,attention_out3,attention_out4,attention_out5,attention_out6,attention_out7,attention_out8,pos_encoding1,pos_encoding2,pos_encoding3,pos_encoding4,pos_encoding5,pos_encoding6,pos_encoding7,pos_encoding8




'''
all_model = model()
input = torch.randn(1,4,1024,1024)
output = all_model(input)
print(output.shape)
'''


all_model = down_model().to(device)
loss_function = nn.MSELoss().to(device)  #2.定义loss
optimizer = torch.optim.Adam(all_model.parameters(),lr=1e-6)  #3.定义优化器

epoch = 3
batch_size = 10
image_size = 268    #【10，3，268，268】


ds = load_dataset("bitmind/ffhq-256",split="train")
preprocess = transforms.Compose(
    [
        transforms.Resize((image_size, image_size)),  # Resize
        transforms.RandomHorizontalFlip(),  # Randomly flip (data augmentation)
        transforms.ToTensor(),  # Convert to tensor (0, 1)
        transforms.Normalize([0.5], [0.5]),  # Map to (-1, 1)
    ]
)
def transform(examples):
    images = [preprocess(image.convert("RGB")) for image in examples["image"]]
    return {"images": images}

ds.set_transform(transform)
dataloader = torch.utils.data.DataLoader(ds,batch_size=batch_size,shuffle=True)


for i in range(epoch):
    for idx, batch_x in enumerate(dataloader):
        images = batch_x["images"].to(device)
        # print(images.shape)  #(4,3,572,572)
        output = all_model(images).to(device)
        loss = loss_function(output, images)
        optimizer.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(all_model.parameters(), 1.)
        optimizer.step()
        print("epoch:", i, "loss:", loss.item())
        wandb.log({'epoch': i,"batch:": idx,'loss':loss})

#torch.save(model.state_dict(), 'model_weights.pth')