Add device set.

This commit is contained in:
梁鸿 2024-08-18 23:47:47 +08:00
parent f2ee49a639
commit 4a8846390b
10 changed files with 13 additions and 11 deletions

Binary file not shown.

Before

Width:  |  Height:  |  Size: 2.8 KiB

After

Width:  |  Height:  |  Size: 2.8 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 191 B

After

Width:  |  Height:  |  Size: 191 B

Binary file not shown.

Before

Width:  |  Height:  |  Size: 187 B

After

Width:  |  Height:  |  Size: 187 B

Binary file not shown.

Before

Width:  |  Height:  |  Size: 1018 B

After

Width:  |  Height:  |  Size: 1020 B

Binary file not shown.

Before

Width:  |  Height:  |  Size: 2.0 KiB

After

Width:  |  Height:  |  Size: 2.0 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 1.8 KiB

After

Width:  |  Height:  |  Size: 1.8 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 84 B

After

Width:  |  Height:  |  Size: 84 B

Binary file not shown.

Before

Width:  |  Height:  |  Size: 2.4 KiB

After

Width:  |  Height:  |  Size: 2.4 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 1.1 KiB

After

Width:  |  Height:  |  Size: 1.1 KiB

View File

@ -13,8 +13,8 @@ seed = 4321
torch.manual_seed(seed) torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed) torch.cuda.manual_seed_all(seed)
# Device configuration
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")
# Hyper-parameters # Hyper-parameters
num_epochs = 1 num_epochs = 1
@ -51,25 +51,25 @@ 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]), "input_image.png") show.DumpTensorToImage(x.view(-1, x.shape[2], x.shape[3]).cpu(), "input_image.png")
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]), "conv1_weight.png") show.DumpTensorToImage(w.view(-1, w.shape[2], w.shape[3]).cpu(), "conv1_weight.png")
show.DumpTensorToImage(x.view(-1, x.shape[2], x.shape[3]), "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
show.DumpTensorToImage(w.view(-1, w.shape[2], w.shape[3]), "conv2_weight.png") show.DumpTensorToImage(w.view(-1, w.shape[2], w.shape[3]).cpu(), "conv2_weight.png")
show.DumpTensorToImage(x.view(-1, x.shape[2], x.shape[3]), "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 * 5 * 5)
x = self.fc1(x) x = self.fc1(x)
show.DumpTensorToImage(self.fc1.weight.view(-1, 10, 10).permute(2, 0, 1), "fc_weight.png") show.DumpTensorToImage(self.fc1.weight.view(-1, 10, 10).permute(2, 0, 1).cpu(), "fc_weight.png")
show.DumpTensorToImage(x.view(-1), "fc_output.png") show.DumpTensorToImage(x.view(-1).cpu(), "fc_output.png")
criterion = nn.CrossEntropyLoss() criterion = nn.CrossEntropyLoss()
loss = criterion(x, label) loss = criterion(x, label)
@ -77,10 +77,10 @@ class ConvNet(nn.Module):
loss.backward() loss.backward()
w = self.conv1.weight.grad w = self.conv1.weight.grad
show.DumpTensorToImage(w.view(-1, w.shape[2], w.shape[3]), "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]), "conv2_weight_grad.png") show.DumpTensorToImage(w.view(-1, w.shape[2], w.shape[3]).cpu(), "conv2_weight_grad.png")
show.DumpTensorToImage(self.fc1.weight.grad.view(-1, 10, 10).permute(2, 0, 1), "fc_weight_grad.png") show.DumpTensorToImage(self.fc1.weight.grad.view(-1, 10, 10).permute(2, 0, 1).cpu(), "fc_weight_grad.png")
model = ConvNet().to(device) model = ConvNet().to(device)
@ -109,6 +109,8 @@ for epoch in range(num_epochs):
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}")
for images, labels in test_loader: for images, labels in test_loader:
images = images.to(device)
labels = labels.to(device)
model.printFector(images, labels) model.printFector(images, labels)
break break