完成实验四

This commit is contained in:
2024-01-12 02:27:03 +08:00
parent 7fbb893223
commit feca24347a
15 changed files with 4588 additions and 7 deletions

73
Lab4/code/1.1.py Normal file
View File

@@ -0,0 +1,73 @@
from utils import *
import ipdb
class My_Conv2d(nn.Module):
def __init__(self, in_channels:int, out_channels:int, kernel_size:int, padding:int=0, bias=True):
super(My_Conv2d, self).__init__()
self.has_bias = bias
self.in_channels = in_channels
self.out_channels = out_channels
self.kernel_size = kernel_size
self.padding = padding
self.weight = nn.Parameter(torch.Tensor(out_channels, in_channels, kernel_size, kernel_size))
nn.init.xavier_uniform_(self.weight)
if self.has_bias:
self.bias = nn.Parameter(torch.zeros(out_channels, requires_grad=True, dtype=torch.float32))
def forward(self, x):
batch_size, _, input_height, input_width = x.shape
if self.padding > 0:
x = F.pad(x, (self.padding, self.padding, self.padding, self.padding))
x = F.unfold(x, kernel_size=self.kernel_size)
x = x.permute(0, 2, 1).contiguous()
weight_unfold = self.weight.view(self.out_channels, -1).t()
x = torch.matmul(x, weight_unfold)
if self.has_bias:
x += self.bias
output_height = input_height + 2 * self.padding - self.kernel_size + 1
output_width = input_width + 2 * self.padding - self.kernel_size + 1
x = x.view(batch_size, output_height, output_width, self.out_channels).permute(0, 3, 1, 2).contiguous()
return x
class Model_Vehicle_CLS_1_1(nn.Module):
def __init__(self, num_classes=3):
super(Model_Vehicle_CLS_1_1, self).__init__()
self.conv1 = My_Conv2d(in_channels=3, out_channels=128, kernel_size=3, padding=1, bias=False)
self.bn1 = nn.BatchNorm2d(128)
self.conv2 = My_Conv2d(in_channels=128, out_channels=512, kernel_size=3, padding=1, bias=False)
self.bn2 = nn.BatchNorm2d(512)
self.fc = nn.Linear(in_features=512, out_features=num_classes)
def forward(self, x):
x = F.relu(self.bn1(self.conv1(x)))
x = F.relu(self.bn2(self.conv2(x)))
x = F.avg_pool2d(x, 32)
x = x.view(x.size(0), -1)
x = self.fc(x)
return x
class Model_Haze_Removal_1_1(nn.Module):
def __init__(self):
super(Model_Haze_Removal_1_1, self).__init__()
self.conv1 = My_Conv2d(in_channels=3, out_channels=16, kernel_size=3, padding=1, bias=False)
self.bn1 = nn.BatchNorm2d(16)
self.conv2 = My_Conv2d(in_channels=16, out_channels=48, kernel_size=5, padding=2, bias=False)
self.bn2 = nn.BatchNorm2d(48)
self.conv3 = My_Conv2d(in_channels=48, out_channels=3, kernel_size=3, padding=1)
def forward(self, x):
x = F.relu(self.bn1(self.conv1(x)))
x = F.relu(self.bn2(self.conv2(x)))
x = self.conv3(x)
return x
if __name__ == "__main__":
model = Model_Vehicle_CLS_1_1()
train_Vehicle_CLS(model=model, learning_rate=4e-4, batch_size=256)
model = Model_Haze_Removal_1_1()
train_Haze_Removal(model=model, learning_rate=5e-3, batch_size=16)

43
Lab4/code/1.2.py Normal file
View File

@@ -0,0 +1,43 @@
from utils import *
import ipdb
class Model_Vehicle_CLS_1_2(nn.Module):
def __init__(self, num_classes=3):
super(Model_Vehicle_CLS_1_2, self).__init__()
self.conv1 = nn.Conv2d(in_channels=3, out_channels=128, kernel_size=3, padding=1, bias=False)
self.bn1 = nn.BatchNorm2d(128)
self.conv2 = nn.Conv2d(in_channels=128, out_channels=512, kernel_size=3, padding=1, bias=False)
self.bn2 = nn.BatchNorm2d(512)
self.fc = nn.Linear(in_features=512, out_features=num_classes)
def forward(self, x):
x = F.relu(self.bn1(self.conv1(x)))
x = F.relu(self.bn2(self.conv2(x)))
x = F.avg_pool2d(x, 32)
x = x.view(x.size(0), -1)
x = self.fc(x)
return x
class Model_Haze_Removal_1_2(nn.Module):
def __init__(self):
super(Model_Haze_Removal_1_2, self).__init__()
self.conv1 = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, padding=1, bias=False)
self.bn1 = nn.BatchNorm2d(16)
self.conv2 = nn.Conv2d(in_channels=16, out_channels=48, kernel_size=5, padding=2, bias=False)
self.bn2 = nn.BatchNorm2d(48)
self.conv3 = nn.Conv2d(in_channels=48, out_channels=3, kernel_size=3, padding=1)
def forward(self, x):
x = F.relu(self.bn1(self.conv1(x)))
x = F.relu(self.bn2(self.conv2(x)))
x = self.conv3(x)
return x
if __name__ == "__main__":
model = Model_Vehicle_CLS_1_2()
train_Vehicle_CLS(model=model, learning_rate=4e-4, batch_size=64)
model = Model_Haze_Removal_1_2()
train_Haze_Removal(model=model, learning_rate=5e-3, batch_size=16)

215
Lab4/code/1.3.py Normal file
View File

@@ -0,0 +1,215 @@
from utils import *
class Model_Vehicle_CLS_1_3_1(nn.Module):
def __init__(self, num_classes=3):
super(Model_Vehicle_CLS_1_3_1, self).__init__()
self.conv1 = nn.Sequential(
nn.Conv2d(in_channels=3, out_channels=512, kernel_size=3, padding=1, bias=False),
nn.BatchNorm2d(512),
)
self.fc = nn.Linear(in_features=512, out_features=num_classes)
def forward(self, x):
x = F.relu(self.conv1(x))
x = F.avg_pool2d(x, 32)
x = x.view(x.size(0), -1)
x = self.fc(x)
return x
class Model_Vehicle_CLS_1_3_2(nn.Module):
def __init__(self, num_classes=3):
super(Model_Vehicle_CLS_1_3_2, self).__init__()
self.conv1 = nn.Sequential(
nn.Conv2d(in_channels=3, out_channels=128, kernel_size=3, padding=1, bias=False),
nn.BatchNorm2d(128),
)
self.conv2 = nn.Sequential(
nn.Conv2d(in_channels=128, out_channels=512, kernel_size=3, padding=1, bias=False),
nn.BatchNorm2d(512),
)
self.fc = nn.Linear(in_features=512, out_features=num_classes, bias=False)
def forward(self, x):
x = F.relu(self.conv1(x))
x = F.relu(self.conv2(x))
x = F.avg_pool2d(x, 32)
x = x.view(x.size(0), -1)
x = self.fc(x)
return x
class Model_Vehicle_CLS_1_3_3(nn.Module):
def __init__(self, num_classes=3):
super(Model_Vehicle_CLS_1_3_3, self).__init__()
self.conv1 = nn.Sequential(
nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3, padding=1, bias=False),
nn.BatchNorm2d(64),
)
self.conv2 = nn.Sequential(
nn.Conv2d(in_channels=64, out_channels=256, kernel_size=3, padding=1, bias=False),
nn.BatchNorm2d(256),
)
self.conv3 = nn.Sequential(
nn.Conv2d(in_channels=256, out_channels=512, kernel_size=3, padding=1, bias=False),
nn.BatchNorm2d(512),
)
self.fc = nn.Linear(in_features=512, out_features=num_classes)
def forward(self, x):
x = F.relu(self.conv1(x))
x = F.relu(self.conv2(x))
x = F.relu(self.conv3(x))
x = F.avg_pool2d(x, 32)
x = x.view(x.size(0), -1)
x = self.fc(x)
return x
class Model_Vehicle_CLS_1_3_4(nn.Module):
def __init__(self, num_classes=3):
super(Model_Vehicle_CLS_1_3_4, self).__init__()
self.conv1 = nn.Sequential(
nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3, padding=1, bias=False),
nn.BatchNorm2d(64),
)
self.conv2 = nn.Sequential(
nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, padding=1, bias=False),
nn.BatchNorm2d(128),
)
self.conv3 = nn.Sequential(
nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, padding=1, bias=False),
nn.BatchNorm2d(256),
)
self.conv4 = nn.Sequential(
nn.Conv2d(in_channels=256, out_channels=512, kernel_size=3, padding=1, bias=False),
nn.BatchNorm2d(512),
)
self.fc = nn.Linear(in_features=512, out_features=num_classes)
def forward(self, x):
x = F.relu(self.conv1(x))
x = F.relu(self.conv2(x))
x = F.relu(self.conv3(x))
x = F.relu(self.conv4(x))
x = F.avg_pool2d(x, 32)
x = x.view(x.size(0), -1)
x = self.fc(x)
return x
class Model_Haze_Removal_1_3_1(nn.Module):
def __init__(self):
super(Model_Haze_Removal_1_3_1, self).__init__()
self.conv1 = nn.Sequential(
nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, padding=1, bias=False),
nn.BatchNorm2d(16),
)
self.conv2 = nn.Sequential(
nn.Conv2d(in_channels=16, out_channels=48, kernel_size=3, padding=1, bias=False),
nn.BatchNorm2d(48),
)
self.conv3 = nn.Conv2d(in_channels=48, out_channels=3, kernel_size=3, padding=1)
def forward(self, x):
x = F.relu(self.conv1(x))
x = F.relu(self.conv2(x))
x = self.conv3(x)
return x
class Model_Haze_Removal_1_3_2(nn.Module):
def __init__(self):
super(Model_Haze_Removal_1_3_2, self).__init__()
self.conv1 = nn.Sequential(
nn.Conv2d(in_channels=3, out_channels=16, kernel_size=5, padding=2, bias=False),
nn.BatchNorm2d(16),
)
self.conv2 = nn.Sequential(
nn.Conv2d(in_channels=16, out_channels=48, kernel_size=5, padding=2, bias=False),
nn.BatchNorm2d(48),
)
self.conv3 = nn.Conv2d(in_channels=48, out_channels=3, kernel_size=5, padding=2)
def forward(self, x):
x = F.relu(self.conv1(x))
x = F.relu(self.conv2(x))
x = self.conv3(x)
return x
class Model_Haze_Removal_1_3_3(nn.Module):
def __init__(self):
super(Model_Haze_Removal_1_3_3, self).__init__()
self.conv1 = nn.Sequential(
nn.Conv2d(in_channels=3, out_channels=16, kernel_size=7, padding=3, bias=False),
nn.BatchNorm2d(16),
nn.ReLU(inplace=True),
)
self.conv2 = nn.Sequential(
nn.Conv2d(in_channels=16, out_channels=48, kernel_size=7, padding=3, bias=False),
nn.BatchNorm2d(48),
nn.ReLU(inplace=True),
)
self.conv3 = nn.Conv2d(in_channels=48, out_channels=3, kernel_size=7, padding=3)
def forward(self, x):
x = F.relu(self.conv1(x))
x = F.relu(self.conv2(x))
x = self.conv3(x)
return x
class Model_Haze_Removal_1_3_4(nn.Module):
def __init__(self):
super(Model_Haze_Removal_1_3_4, self).__init__()
self.conv1 = nn.Sequential(
nn.Conv2d(in_channels=3, out_channels=16, kernel_size=9, padding=4, bias=False),
nn.BatchNorm2d(16),
nn.ReLU(inplace=True),
)
self.conv2 = nn.Sequential(
nn.Conv2d(in_channels=16, out_channels=48, kernel_size=9, padding=4, bias=False),
nn.BatchNorm2d(48),
nn.ReLU(inplace=True),
)
self.conv3 = nn.Conv2d(in_channels=48, out_channels=3, kernel_size=9, padding=4)
def forward(self, x):
x = F.relu(self.conv1(x))
x = F.relu(self.conv2(x))
x = self.conv3(x)
return x
if __name__ == "__main__":
num_epochs = 61
learning_rate = 2e-4
batch_size = 256
models = [
Model_Vehicle_CLS_1_3_1,
Model_Vehicle_CLS_1_3_2,
Model_Vehicle_CLS_1_3_3,
Model_Vehicle_CLS_1_3_4,
]
for i in range(4):
model = models[i]()
print(f"卷积层层数={i + 1}")
train_loss, test_acc = train_Vehicle_CLS(model=model, learning_rate=learning_rate,
batch_size=batch_size, num_epochs=num_epochs)
print()
num_epochs = 61
learning_rate = 8e-3
batch_size = 64
models = [
Model_Haze_Removal_1_3_1,
Model_Haze_Removal_1_3_2,
Model_Haze_Removal_1_3_3,
Model_Haze_Removal_1_3_4,
]
for i in range(4):
model = models[i]()
print(f"卷积核大小={3 + 2 * i}")
train_loss, test_loss = train_Haze_Removal(model=model, learning_rate=learning_rate, batch_size=batch_size, num_epochs=num_epochs)
print()

168
Lab4/code/2.py Normal file
View File

@@ -0,0 +1,168 @@
from utils import *
class Model_Vehicle_CLS_2_1(nn.Module):
def __init__(self, num_classes=3):
super(Model_Vehicle_CLS_2_1, self).__init__()
self.conv1 = nn.Sequential(
nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, padding=1, dilation=1, bias=False),
nn.BatchNorm2d(16),
nn.ReLU(inplace=True),
nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, padding=1, dilation=1, bias=False),
nn.BatchNorm2d(32),
nn.ReLU(inplace=True),
nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, padding=1, dilation=1, bias=False),
nn.BatchNorm2d(64),
nn.ReLU(inplace=True),
)
self.conv2 = nn.Sequential(
nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, padding=1, dilation=1, bias=False),
nn.BatchNorm2d(128),
nn.ReLU(inplace=True),
nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, padding=1, dilation=1, bias=False),
nn.BatchNorm2d(256),
nn.ReLU(inplace=True),
nn.Conv2d(in_channels=256, out_channels=512, kernel_size=3, padding=1, dilation=1, bias=False),
nn.BatchNorm2d(512),
nn.ReLU(inplace=True),
)
self.fc = nn.Linear(in_features=512, out_features=num_classes)
def forward(self, x):
x = self.conv1(x)
x = self.conv2(x)
x = F.avg_pool2d(x, 32)
x = x.view(x.size(0), -1)
x = self.fc(x)
return x
class Model_Vehicle_CLS_2_2(nn.Module):
def __init__(self, num_classes=3):
super(Model_Vehicle_CLS_2_2, self).__init__()
self.conv1 = nn.Sequential(
nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, padding=1, dilation=1, bias=False),
nn.BatchNorm2d(16),
nn.ReLU(inplace=True),
nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, padding=2, dilation=2, bias=False),
nn.BatchNorm2d(32),
nn.ReLU(inplace=True),
nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, padding=5, dilation=5, bias=False),
nn.BatchNorm2d(64),
nn.ReLU(inplace=True),
)
self.conv2 = nn.Sequential(
nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, padding=1, dilation=1, bias=False),
nn.BatchNorm2d(128),
nn.ReLU(inplace=True),
nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, padding=2, dilation=2, bias=False),
nn.BatchNorm2d(256),
nn.ReLU(inplace=True),
nn.Conv2d(in_channels=256, out_channels=512, kernel_size=3, padding=5, dilation=5, bias=False),
nn.BatchNorm2d(512),
nn.ReLU(inplace=True),
)
self.fc = nn.Linear(in_features=512, out_features=num_classes)
def forward(self, x):
x = self.conv1(x)
x = self.conv2(x)
x = F.avg_pool2d(x, 32)
x = x.view(x.size(0), -1)
x = self.fc(x)
return x
class Model_Vehicle_CLS_2_3(nn.Module):
def __init__(self, num_classes=3):
super(Model_Vehicle_CLS_2_3, self).__init__()
self.conv1 = nn.Sequential(
nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, padding=1, dilation=1, bias=False),
nn.BatchNorm2d(16),
nn.ReLU(inplace=True),
nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, padding=3, dilation=3, bias=False),
nn.BatchNorm2d(32),
nn.ReLU(inplace=True),
nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, padding=5, dilation=5, bias=False),
nn.BatchNorm2d(64),
nn.ReLU(inplace=True),
)
self.conv2 = nn.Sequential(
nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, padding=1, dilation=1, bias=False),
nn.BatchNorm2d(128),
nn.ReLU(inplace=True),
nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, padding=3, dilation=3, bias=False),
nn.BatchNorm2d(256),
nn.ReLU(inplace=True),
nn.Conv2d(in_channels=256, out_channels=512, kernel_size=3, padding=5, dilation=5, bias=False),
nn.BatchNorm2d(512),
nn.ReLU(inplace=True),
)
self.fc = nn.Linear(in_features=512, out_features=num_classes)
def forward(self, x):
x = self.conv1(x)
x = self.conv2(x)
x = F.avg_pool2d(x, 32)
x = x.view(x.size(0), -1)
x = self.fc(x)
return x
class Model_Vehicle_CLS_2_4(nn.Module):
def __init__(self, num_classes=3):
super(Model_Vehicle_CLS_2_4, self).__init__()
self.conv1 = nn.Sequential(
nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, padding=1, dilation=1, bias=False),
nn.BatchNorm2d(16),
nn.ReLU(inplace=True),
nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, padding=3, dilation=3, bias=False),
nn.BatchNorm2d(32),
nn.ReLU(inplace=True),
nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, padding=7, dilation=7, bias=False),
nn.BatchNorm2d(64),
nn.ReLU(inplace=True),
)
self.conv2 = nn.Sequential(
nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, padding=1, dilation=1, bias=False),
nn.BatchNorm2d(128),
nn.ReLU(inplace=True),
nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, padding=3, dilation=3, bias=False),
nn.BatchNorm2d(256),
nn.ReLU(inplace=True),
nn.Conv2d(in_channels=256, out_channels=512, kernel_size=3, padding=7, dilation=7, bias=False),
nn.BatchNorm2d(512),
nn.ReLU(inplace=True),
)
self.fc = nn.Linear(in_features=512, out_features=num_classes)
def forward(self, x):
x = self.conv1(x)
x = self.conv2(x)
x = F.avg_pool2d(x, 32)
x = x.view(x.size(0), -1)
x = self.fc(x)
return x
if __name__ == "__main__":
num_epochs = 61
learning_rate = 1e-4
batch_size = 256
dilations = ["[[1, 1, 1], [1, 1, 1]]",
"[[1, 2, 5], [1, 2, 5]]",
"[[1, 3, 5], [1, 3, 5]]",
"[[1, 3, 7], [1, 3, 7]]"]
models = [
Model_Vehicle_CLS_2_1,
Model_Vehicle_CLS_2_2,
Model_Vehicle_CLS_2_3,
Model_Vehicle_CLS_2_4,
]
for i in range(4):
model = models[i]()
print("Dilation: " + dilations[i])
train_loss, test_acc = train_Vehicle_CLS(model=model, learning_rate=learning_rate,
batch_size=batch_size, num_epochs=num_epochs)
print()

63
Lab4/code/3.py Normal file
View File

@@ -0,0 +1,63 @@
from utils import *
class BasicResidualBlock(nn.Module):
def __init__(self, in_channels, out_channels, stride=1):
super(BasicResidualBlock, self).__init__()
self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False)
self.bn1 = nn.BatchNorm2d(out_channels)
self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False)
self.bn2 = nn.BatchNorm2d(out_channels)
self.shortcut = nn.Sequential()
if stride != 1 or in_channels != out_channels:
self.shortcut = nn.Sequential(
nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride, bias=False),
nn.BatchNorm2d(out_channels)
)
def forward(self, x):
out = F.relu(self.bn1(self.conv1(x)))
out = self.bn2(self.conv2(out))
out += self.shortcut(x)
out = F.relu(out)
return out
class Model_Vehicle_CLS_3(nn.Module):
def __init__(self, num_classes=3):
super(Model_Vehicle_CLS_3, self).__init__()
self.conv1 = nn.Sequential(
nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3, padding=1, bias=False),
nn.BatchNorm2d(64),
)
self.conv2 = BasicResidualBlock(in_channels=64, out_channels=64)
self.conv3 = BasicResidualBlock(in_channels=64, out_channels=64)
self.conv4 = BasicResidualBlock(in_channels=64, out_channels=128, stride=2)
self.conv5 = BasicResidualBlock(in_channels=128, out_channels=128)
self.conv6 = BasicResidualBlock(in_channels=128, out_channels=256, stride=2)
self.conv7 = BasicResidualBlock(in_channels=256, out_channels=256)
self.conv8 = BasicResidualBlock(in_channels=256, out_channels=512, stride=2)
self.conv9 = BasicResidualBlock(in_channels=512, out_channels=512)
self.fc = nn.Linear(in_features=512, out_features=num_classes)
def forward(self, x):
x = F.relu(self.conv1(x))
x = self.conv2(x)
x = self.conv3(x)
x = self.conv4(x)
x = self.conv5(x)
x = self.conv6(x)
x = self.conv7(x)
x = self.conv8(x)
x = self.conv9(x)
x = F.avg_pool2d(x, 4)
x = x.view(x.size(0), -1)
x = self.fc(x)
return x
if __name__ == "__main__":
num_epochs = 61
learning_rate = 15e-5
batch_size = 512
model = Model_Vehicle_CLS_3(num_classes=3)
train_loss, test_acc = train_Vehicle_CLS(model=model, learning_rate=learning_rate, batch_size=batch_size, num_epochs=num_epochs)

197
Lab4/code/utils.py Normal file
View File

@@ -0,0 +1,197 @@
import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset
from torch import nn
from torchvision import transforms
from tqdm import tqdm
import os
import time
from PIL import Image
import pandas as pd
class Vehicle(Dataset):
def __init__(self, root:str="../dataset", train:bool=True, transform=None):
root = os.path.join(root, "Vehicles")
csv_file = os.path.join(root, "train.csv" if train else "test.csv")
self.data = pd.read_csv(csv_file).to_numpy().tolist()
self.root = root
self.transform = transform
def __len__(self):
return len(self.data)
def __getitem__(self, index):
img_name, label = self.data[index]
img_path = os.path.join(self.root, img_name)
image = Image.open(img_path)
label = int(label)
if self.transform:
image = self.transform(image)
return image, label
class Haze(Dataset):
def __init__(self, root:str="../dataset", train:bool=True, transform=None):
root = os.path.join(root, "Haze")
split_file = pd.read_csv(os.path.join(root, "split.csv")).to_numpy().tolist()
self.data = list()
for img, is_train in split_file:
if train and int(is_train) == 1:
self.data.append(img)
elif not train and int(is_train) == 0:
self.data.append(img)
self.root = root
self.transform = transform
def __len__(self):
return len(self.data)
def __getitem__(self, index):
img_name = self.data[index]
img_path = os.path.join(self.root, "raw/haze", img_name)
ground_truth_path = os.path.join(self.root, "raw/no_haze", img_name)
image = Image.open(img_path)
ground_truth = Image.open(ground_truth_path)
if self.transform:
image = self.transform(image)
ground_truth = self.transform(ground_truth)
return image, ground_truth
def train_Vehicle_CLS(model:nn.Module, learning_rate=1e-3, batch_size=64, num_epochs=51):
num_classes = 3
device = "cuda:0" if torch.cuda.is_available() else "cpu"
transform = transforms.Compose(
[
transforms.ToTensor(),
transforms.Resize((32, 32), antialias=True),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
]
)
train_dataset = Vehicle(root="../dataset", train=True, transform=transform)
test_dataset = Vehicle(root="../dataset", train=False, transform=transform)
train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True, num_workers=14, pin_memory=True)
test_loader = DataLoader(dataset=test_dataset, batch_size=batch_size, num_workers=14, pin_memory=True)
model = model.to(device)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
train_loss = list()
test_acc = list()
for epoch in range(num_epochs):
model.train()
total_epoch_loss = 0
train_tik = time.time()
for index, (images, targets) in tqdm(enumerate(train_loader), total=len(train_loader)):
optimizer.zero_grad()
images = images.to(device)
targets = targets.to(device)
one_hot_targets = F.one_hot(targets, num_classes=num_classes).to(dtype=torch.float)
outputs = model(images)
loss = criterion(outputs, one_hot_targets)
total_epoch_loss += loss.item()
loss.backward()
optimizer.step()
train_tok = time.time()
model.eval()
with torch.no_grad():
total_epoch_acc = 0
test_tik = time.time()
for index, (image, targets) in tqdm(enumerate(test_loader), total=len(test_loader)):
image = image.to(device)
targets = targets.to(device)
outputs = model(image)
pred = F.softmax(outputs, dim=1)
total_epoch_acc += (pred.argmax(1) == targets).sum().item()
test_tok = time.time()
avg_epoch_acc = total_epoch_acc / len(test_dataset)
print(
f"Epoch [{epoch + 1}/{num_epochs}],",
f"Train Loss: {total_epoch_loss:.10f},",
f"Train Time: {1000 * (train_tok - train_tik):.2f}ms,",
f"Test Acc: {avg_epoch_acc * 100:.3f}%,",
f"Test Time: {1000 * (test_tok - test_tik):.2f}ms"
)
train_loss.append(total_epoch_loss)
test_acc.append(avg_epoch_acc)
print(f"最大显存使用量: {torch.cuda.max_memory_allocated() / (1024 * 1024):.2f}MiB")
torch.cuda.reset_peak_memory_stats()
return train_loss, test_acc
def train_Haze_Removal(model:nn.Module, learning_rate=1e-3, batch_size=64, num_epochs=51):
num_epochs = 50
device = "cuda:0" if torch.cuda.is_available() else "cpu"
transform = transforms.Compose(
[
transforms.ToTensor(),
transforms.Resize((224, 224), antialias=True),
]
)
train_dataset = Haze(root="../dataset", train=True, transform=transform)
test_dataset = Haze(root="../dataset", train=False, transform=transform)
train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True, num_workers=14, pin_memory=True)
test_loader = DataLoader(dataset=test_dataset, batch_size=batch_size, num_workers=14, pin_memory=True)
model = model.to(device)
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
train_loss = list()
test_loss = list()
for epoch in range(num_epochs):
model.train()
total_epoch_train_loss = 0
train_tik = time.time()
for index, (images, ground_truth) in tqdm(enumerate(train_loader), total=len(train_loader)):
optimizer.zero_grad()
images = images.to(device)
ground_truth = ground_truth.to(device)
outputs = model(images)
loss = criterion(outputs, ground_truth)
total_epoch_train_loss += loss.item()
loss.backward()
optimizer.step()
train_tok = time.time()
model.eval()
with torch.no_grad():
total_epoch_test_loss = 0
test_tik = time.time()
for index, (image, ground_truth) in tqdm(enumerate(test_loader), total=len(test_loader)):
image = image.to(device)
ground_truth = ground_truth.to(device)
outputs = model(image)
loss = criterion(outputs, ground_truth)
total_epoch_test_loss += loss.item()
test_tok = time.time()
print(
f"Epoch [{epoch + 1}/{num_epochs}],",
f"Train Loss: {total_epoch_train_loss:.10f},",
f"Train Time: {1000 * (train_tok - train_tik):.2f}ms,",
f"Test Loss: {total_epoch_test_loss:.10f},",
f"Test Time: {1000 * (test_tok - test_tik):.2f}ms"
)
train_loss.append(total_epoch_train_loss)
test_loss.append(total_epoch_test_loss)
print(f"最大显存使用量: {torch.cuda.max_memory_allocated() / (1024 * 1024):.2f}MiB")
torch.cuda.reset_peak_memory_stats()
return train_loss, test_loss