Refine to dump heat image.

This commit is contained in:
Colin 2024-08-29 16:47:06 +08:00
parent 4a8846390b
commit c395fa7baa
12 changed files with 49 additions and 27 deletions

View File

@ -8,7 +8,7 @@ import numpy as np
import os import os
def DumpTensorToImage(tensor, name, forceSquare=True, scale=1.0, AutoContrast=True, GridValue=0): def DumpTensorToImage(tensor, name, forceSquare=True, scale=1.0, Contrast=None, GridValue=None):
if len(tensor.shape) != 2 and len(tensor.shape) != 1 and len(tensor.shape) != 3: if len(tensor.shape) != 2 and len(tensor.shape) != 1 and len(tensor.shape) != 3:
raise ("Error input dims") raise ("Error input dims")
if ("." not in name) or (name.split(".")[-1] not in {"jpg", "png", "bmp"}): if ("." not in name) or (name.split(".")[-1] not in {"jpg", "png", "bmp"}):
@ -17,27 +17,36 @@ def DumpTensorToImage(tensor, name, forceSquare=True, scale=1.0, AutoContrast=Tr
if len(tensor.shape) == 3: if len(tensor.shape) == 3:
channel = tensor.shape[0] channel = tensor.shape[0]
x = math.ceil((channel) ** 0.5) x = math.ceil((channel) ** 0.5)
tensor = F.pad(tensor, (0, 0, 0, 0, 0, x * x - channel), mode="constant", value=0) calc = tensor.reshape((channel, tensor.shape[1] * tensor.shape[2]))
calc = tensor.reshape((x * x, tensor.shape[1] * tensor.shape[2])) if not Contrast:
if AutoContrast:
tensormax = calc.max(1)[0] tensormax = calc.max(1)[0]
tensormin = calc.min(1)[0] tensormin = calc.min(1)[0]
calc = calc.transpose(1, 0) else:
calc = ((calc - tensormin) / (tensormax - tensormin)) * 255 tensormax = Contrast[1]
calc = calc.transpose(1, 0) tensormin = Contrast[0]
calc = calc.transpose(1, 0)
calc = ((calc - tensormin) / (tensormax - tensormin)) * 255.0
calc = calc.transpose(1, 0)
calc = calc.reshape((channel, tensor.shape[1], tensor.shape[2]))
if not GridValue:
GridValue = 128.0
calc = F.pad(calc, (0, 0, 0, 0, 0, x * x - channel), mode="constant", value=GridValue)
calc = calc.reshape((x, x, tensor.shape[1], tensor.shape[2])) calc = calc.reshape((x, x, tensor.shape[1], tensor.shape[2]))
calc = F.pad(calc, (0, 1, 0, 1, 0, 0), mode="constant", value=GridValue) calc = F.pad(calc, (0, 1, 0, 1, 0, 0), mode="constant", value=GridValue)
tensor = calc.permute((0, 2, 1, 3)) tensor = calc.permute((0, 2, 1, 3))
tensor = tensor.reshape((x * tensor.shape[1], x * tensor.shape[3])) tensor = tensor.reshape((x * tensor.shape[1], x * tensor.shape[3]))
DumpTensorToImage(tensor, name, forceSquare=False, scale=scale, AutoContrast=False) DumpTensorToImage(tensor, name, forceSquare=False, scale=scale, Contrast=[0.0, 255.0], GridValue=GridValue)
return return
tensor = tensor.float() tensor = tensor.float()
if AutoContrast: if not Contrast:
maxv = torch.max(tensor) maxv = torch.max(tensor)
minv = torch.min(tensor) minv = torch.min(tensor)
tensor = ((tensor - minv) / (maxv - minv)) * 255 else:
img = tensor.byte().cpu().numpy() maxv = Contrast[1]
minv = Contrast[0]
tensor = ((tensor - minv) / (maxv - minv)) * 255.0
img = tensor.detach().cpu().numpy()
srp = img.shape srp = img.shape
if len(srp) == 1: # 1D的数据自动折叠成2D图像 if len(srp) == 1: # 1D的数据自动折叠成2D图像
@ -51,13 +60,21 @@ def DumpTensorToImage(tensor, name, forceSquare=True, scale=1.0, AutoContrast=Tr
if scale != 1.0: if scale != 1.0:
img = cv2.resize(img, [int(srp[0] * scale), int(srp[1] * scale)]) img = cv2.resize(img, [int(srp[0] * scale), int(srp[1] * scale)])
srp = img.shape srp = img.shape
cv2.imwrite(name, img)
img = img * (-1)
img = img + 255
img[img < 0] = 0
img = np.nan_to_num(img, nan=0.0)
img[img > 255] = 255
imgs = img.astype(np.uint8)
imgs = cv2.applyColorMap(imgs, cv2.COLORMAP_JET)
cv2.imwrite(name, imgs)
def DumpTensorToLog(tensor, name="log"): def DumpTensorToLog(tensor, name="log"):
shape = tensor.shape shape = tensor.shape
f = open(name, "w") f = open(name, "w")
data = tensor.reshape([-1]).float().cpu().numpy().tolist() data = tensor.reshape([-1]).float().cpu().detach().numpy().tolist()
for d in data: for d in data:
f.writelines("%s" % d + os.linesep) f.writelines("%s" % d + os.linesep)
f.close() f.close()

Binary file not shown.

Before

Width:  |  Height:  |  Size: 2.8 KiB

After

Width:  |  Height:  |  Size: 6.4 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 191 B

After

Width:  |  Height:  |  Size: 648 B

Binary file not shown.

Before

Width:  |  Height:  |  Size: 187 B

After

Width:  |  Height:  |  Size: 698 B

Binary file not shown.

Before

Width:  |  Height:  |  Size: 1020 B

After

Width:  |  Height:  |  Size: 1.6 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 2.0 KiB

After

Width:  |  Height:  |  Size: 4.8 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 1.8 KiB

After

Width:  |  Height:  |  Size: 4.2 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 84 B

After

Width:  |  Height:  |  Size: 109 B

Binary file not shown.

Before

Width:  |  Height:  |  Size: 2.4 KiB

After

Width:  |  Height:  |  Size: 3.4 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 1.1 KiB

After

Width:  |  Height:  |  Size: 2.0 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 320 B

After

Width:  |  Height:  |  Size: 544 B

View File

@ -14,11 +14,11 @@ torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed) torch.cuda.manual_seed_all(seed)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
device = torch.device("mps") # device = torch.device("mps")
# Hyper-parameters # Hyper-parameters
num_epochs = 1 num_epochs = 1
batch_size = 1 batch_size = 256
learning_rate = 0.001 learning_rate = 0.001
transform = transforms.Compose([transforms.ToTensor()]) transform = transforms.Compose([transforms.ToTensor()])
@ -32,17 +32,17 @@ test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, s
class ConvNet(nn.Module): class ConvNet(nn.Module):
def __init__(self): def __init__(self):
super(ConvNet, self).__init__() super(ConvNet, self).__init__()
self.conv1 = nn.Conv2d(1, 8, 3, 1, 1) self.conv1 = nn.Conv2d(1, 8, 5, 1, 0)
self.pool = nn.MaxPool2d(2, 2) self.pool = nn.MaxPool2d(2, 2)
self.conv2 = nn.Conv2d(8, 8, 5) self.conv2 = nn.Conv2d(8, 8, 5, 1, 0)
self.fc1 = nn.Linear(8 * 5 * 5, 10) self.fc1 = nn.Linear(8 * 4 * 4, 10)
# self.fc2 = nn.Linear(120, 84) # self.fc2 = nn.Linear(120, 84)
# self.fc3 = nn.Linear(84, 10) # self.fc3 = nn.Linear(84, 10)
def forward(self, x): def forward(self, x):
x = self.pool(F.relu(self.conv1(x))) x = self.pool(F.relu(self.conv1(x)))
x = self.pool(F.relu(self.conv2(x))) x = self.pool(F.relu(self.conv2(x)))
x = x.view(-1, 8 * 5 * 5) x = x.view(-1, 8 * 4 * 4)
# x = F.relu(self.fc1(x)) # x = F.relu(self.fc1(x))
# x = F.relu(self.fc2(x)) # x = F.relu(self.fc2(x))
# x = self.fc3(x) # x = self.fc3(x)
@ -51,12 +51,16 @@ class ConvNet(nn.Module):
return x return x
def printFector(self, x, label): def printFector(self, x, label):
show.DumpTensorToImage(x.view(-1, x.shape[2], x.shape[3]).cpu(), "input_image.png") show.DumpTensorToImage(x.view(-1, x.shape[2], x.shape[3]), "input_image.png", Contrast=[0, 1.0])
# show.DumpTensorToLog(x, "input_image.log")
x = self.conv1(x) x = self.conv1(x)
w = self.conv1.weight w = self.conv1.weight
show.DumpTensorToImage(w.view(-1, w.shape[2], w.shape[3]).cpu(), "conv1_weight.png") show.DumpTensorToImage(w.view(-1, w.shape[2], w.shape[3]), "conv1_weight.png", Contrast=[-1.0, 1.0])
# show.DumpTensorToLog(w, "conv1_weight.log")
show.DumpTensorToImage(x.view(-1, x.shape[2], x.shape[3]), "conv1_output.png", Contrast=[-1.0, 1.0])
# show.DumpTensorToLog(x, "conv1_output.png")
show.DumpTensorToImage(x.view(-1, x.shape[2], x.shape[3]).cpu(), "conv1_output.png")
x = self.pool(F.relu(x)) x = self.pool(F.relu(x))
x = self.conv2(x) x = self.conv2(x)
w = self.conv2.weight w = self.conv2.weight
@ -64,10 +68,10 @@ class ConvNet(nn.Module):
show.DumpTensorToImage(x.view(-1, x.shape[2], x.shape[3]).cpu(), "conv2_output.png") show.DumpTensorToImage(x.view(-1, x.shape[2], x.shape[3]).cpu(), "conv2_output.png")
x = self.pool(F.relu(x)) x = self.pool(F.relu(x))
x = x.view(-1, 8 * 5 * 5) x = x.view(-1, 8 * 4 * 4)
x = self.fc1(x) x = self.fc1(x)
show.DumpTensorToImage(self.fc1.weight.view(-1, 10, 10).permute(2, 0, 1).cpu(), "fc_weight.png") show.DumpTensorToImage(self.fc1.weight.view(-1, 16, 8).permute(2, 0, 1), "fc_weight.png")
show.DumpTensorToImage(x.view(-1).cpu(), "fc_output.png") show.DumpTensorToImage(x.view(-1).cpu(), "fc_output.png")
@ -79,8 +83,8 @@ class ConvNet(nn.Module):
w = self.conv1.weight.grad w = self.conv1.weight.grad
show.DumpTensorToImage(w.view(-1, w.shape[2], w.shape[3]).cpu(), "conv1_weight_grad.png") show.DumpTensorToImage(w.view(-1, w.shape[2], w.shape[3]).cpu(), "conv1_weight_grad.png")
w = self.conv2.weight.grad w = self.conv2.weight.grad
show.DumpTensorToImage(w.view(-1, w.shape[2], w.shape[3]).cpu(), "conv2_weight_grad.png") show.DumpTensorToImage(w.view(-1, w.shape[2], w.shape[3]), "conv2_weight_grad.png")
show.DumpTensorToImage(self.fc1.weight.grad.view(-1, 10, 10).permute(2, 0, 1).cpu(), "fc_weight_grad.png") show.DumpTensorToImage(self.fc1.weight.grad.view(-1, 16, 8).permute(2, 0, 1), "fc_weight_grad.png")
model = ConvNet().to(device) model = ConvNet().to(device)
@ -105,9 +109,10 @@ for epoch in range(num_epochs):
loss.backward() loss.backward()
optimizer.step() optimizer.step()
if (i + 1) % 2000 == 0: if (i + 1) % 100 == 0:
print(f"Epoch [{epoch+1}/{num_epochs}], Step [{i+1}/{n_total_steps}], Loss: {loss.item():.4f}") print(f"Epoch [{epoch+1}/{num_epochs}], Step [{i+1}/{n_total_steps}], Loss: {loss.item():.4f}")
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=1, shuffle=False)
for images, labels in test_loader: for images, labels in test_loader:
images = images.to(device) images = images.to(device)
labels = labels.to(device) labels = labels.to(device)