mirror of
https://github.com/wassname/CoordConv.git
synced 2026-09-07 16:31:02 +08:00
195 lines
6.9 KiB
Python
195 lines
6.9 KiB
Python
import numpy as np
|
|
import matplotlib.pyplot as plt
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.utils.data as utils
|
|
import torch.nn.functional as F
|
|
|
|
from torch.autograd import Variable
|
|
# from torchsummary import summary
|
|
|
|
datatype = 'uniform'
|
|
assert datatype in ['uniform', 'quadrant']
|
|
|
|
if datatype == 'uniform':
|
|
# Load the one hot datasets
|
|
train_onehot = np.load('data-uniform/train_onehot.npy').astype('float32')
|
|
test_onehot = np.load('data-uniform/test_onehot.npy').astype('float32')
|
|
|
|
# (N, C, H, W) <=== 数据格式
|
|
# make the train and test datasets
|
|
# train
|
|
pos_train = np.where(train_onehot == 1.0)
|
|
X_train = pos_train[2]
|
|
Y_train = pos_train[3]
|
|
train_set = np.zeros((len(X_train), 2, 1, 1), dtype='float32')
|
|
for i, (x, y) in enumerate(zip(X_train, Y_train)):
|
|
train_set[i, 0, 0, 0] = x
|
|
train_set[i, 1, 0, 0] = y
|
|
|
|
# test
|
|
pos_test = np.where(test_onehot == 1.0)
|
|
X_test = pos_test[2]
|
|
Y_test = pos_test[3]
|
|
test_set = np.zeros((len(X_test), 2, 1, 1), dtype='float32')
|
|
for i, (x, y) in enumerate(zip(X_test, Y_test)):
|
|
test_set[i, 0, 0, 0] = x
|
|
test_set[i, 1, 0, 0] = y
|
|
|
|
train_set = np.tile(train_set, [1, 1, 64, 64])
|
|
test_set = np.tile(test_set, [1, 1, 64, 64])
|
|
|
|
# Normalize the datasets
|
|
train_set /= (64. - 1.) # 64x64 grid, 0-based index
|
|
test_set /= (64. - 1.) # 64x64 grid, 0-based index
|
|
|
|
print('Train set : ', train_set.shape, train_set.max(), train_set.min())
|
|
print('Test set : ', test_set.shape, test_set.max(), test_set.min())
|
|
|
|
# Visualize the datasets
|
|
|
|
plt.imshow(np.sum(train_onehot, axis=0)[0, :, :], cmap='gray')
|
|
plt.title('Train One-hot dataset')
|
|
plt.show()
|
|
plt.imshow(np.sum(test_onehot, axis=0)[0, :, :], cmap='gray')
|
|
plt.title('Test One-hot dataset')
|
|
plt.show()
|
|
|
|
else:
|
|
# Load the one hot datasets and the train / test set
|
|
train_set = np.load('data-quadrant/train_set.npy').astype('float32')
|
|
test_set = np.load('data-quadrant/test_set.npy').astype('float32')
|
|
|
|
train_onehot = np.load('data-quadrant/train_onehot.npy').astype('float32')
|
|
test_onehot = np.load('data-quadrant/test_onehot.npy').astype('float32')
|
|
|
|
train_set = np.tile(train_set, [1, 1, 64, 64])
|
|
test_set = np.tile(test_set, [1, 1, 64, 64])
|
|
|
|
# Normalize datasets
|
|
train_set /= train_set.max()
|
|
test_set /= test_set.max()
|
|
|
|
print('Train set : ', train_set.shape, train_set.max(), train_set.min())
|
|
print('Test set : ', test_set.shape, test_set.max(), test_set.min())
|
|
|
|
# Visualize the datasets
|
|
|
|
plt.imshow(np.sum(train_onehot, axis=0)[0, :, :], cmap='gray')
|
|
plt.title('Train One-hot dataset')
|
|
plt.show()
|
|
plt.imshow(np.sum(test_onehot, axis=0)[0, :, :], cmap='gray')
|
|
plt.title('Test One-hot dataset')
|
|
plt.show()
|
|
|
|
# flatten the datasets
|
|
train_onehot = train_onehot.reshape((-1, 64 * 64)).astype('int64')
|
|
test_onehot = test_onehot.reshape((-1, 64 * 64)).astype('int64')
|
|
|
|
# model definition
|
|
|
|
class Net(nn.Module):
|
|
def __init__(self):
|
|
super(Net, self).__init__()
|
|
self.coordconv = CoordConv2d(2, 32, 1, with_r=True)
|
|
self.conv1 = nn.Conv2d(32, 64, 1)
|
|
self.conv2 = nn.Conv2d(64, 64, 1)
|
|
self.conv3 = nn.Conv2d(64, 1, 1)
|
|
self.conv4 = nn.Conv2d( 1, 1, 1)
|
|
|
|
def forward(self, x):
|
|
x = self.coordconv(x)
|
|
x = F.relu(self.conv1(x))
|
|
x = F.relu(self.conv2(x))
|
|
x = F.relu(self.conv3(x))
|
|
x = self.conv4(x)
|
|
x = x.view(-1, 64*64)
|
|
return x
|
|
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
net = Net().to(device)
|
|
|
|
# summary(net, input_size=(2, 64, 64))
|
|
# ----------------------------------------------------------------
|
|
# Layer (type) Output Shape Param #
|
|
# ================================================================
|
|
# AddCoords-1 [-1, 5, 64, 64] 0
|
|
# Conv2d-2 [-1, 32, 64, 64] 192
|
|
# CoordConv2d-3 [-1, 32, 64, 64] 96
|
|
# Conv2d-4 [-1, 64, 64, 64] 2,112
|
|
# Conv2d-5 [-1, 64, 64, 64] 4,160
|
|
# Conv2d-6 [-1, 1, 64, 64] 65
|
|
# Conv2d-7 [-1, 1, 64, 64] 2
|
|
# ================================================================
|
|
# Total params: 6,627
|
|
# Trainable params: 6,627
|
|
# Non-trainable params: 0
|
|
# ----------------------------------------------------------------
|
|
|
|
train_tensor_x = torch.stack([torch.Tensor(i) for i in train_set])
|
|
train_tensor_y = torch.stack([torch.LongTensor(i) for i in train_onehot])
|
|
|
|
train_dataset = utils.TensorDataset(train_tensor_x,train_tensor_y)
|
|
train_dataloader = utils.DataLoader(train_dataset, batch_size=32, shuffle=False)
|
|
|
|
test_tensor_x = torch.stack([torch.Tensor(i) for i in test_set])
|
|
test_tensor_y = torch.stack([torch.LongTensor(i) for i in test_onehot])
|
|
|
|
test_dataset = utils.TensorDataset(test_tensor_x,test_tensor_y)
|
|
test_dataloader = utils.DataLoader(test_dataset, batch_size=32, shuffle=False)
|
|
|
|
optimizer = torch.optim.Adam(net.parameters(), lr=1e-3)
|
|
|
|
def cross_entropy_one_hot(input, target):
|
|
_, labels = target.max(dim=1)
|
|
return nn.CrossEntropyLoss()(input, labels)
|
|
|
|
criterion = cross_entropy_one_hot
|
|
|
|
epochs = 10
|
|
|
|
def train(epoch, net, train_dataloader, optimizer, criterion, device):
|
|
net.train()
|
|
iters = 0
|
|
for batch_idx, (data, target) in enumerate(train_dataloader):
|
|
data, target = Variable(data), Variable(target)
|
|
data, target = data.to(device), target.to(device)
|
|
optimizer.zero_grad()
|
|
output = net(data)
|
|
loss = criterion(output, target)
|
|
loss.backward()
|
|
optimizer.step()
|
|
iters += len(data)
|
|
print('Train Epoch: {} [{}/{} ({:.0f}%)] Loss: {:.6f}'.format(
|
|
epoch, iters, len(train_dataloader.dataset),
|
|
100. * (batch_idx + 1) / len(train_dataloader), loss.data.item()), end='\r', flush=True)
|
|
print("")
|
|
|
|
|
|
for epoch in range(1, epochs + 1):
|
|
train(epoch, net, train_dataloader, optimizer, criterion, device)
|
|
|
|
|
|
def test(net, test_loader, optimizer, criterion, device):
|
|
net.eval()
|
|
test_loss = 0
|
|
correct = 0
|
|
for data, target in test_loader:
|
|
with torch.no_grad():
|
|
data, target = Variable(data), Variable(target)
|
|
data, target = data.to(device), target.to(device)
|
|
output = net(data)
|
|
test_loss += criterion(output, target).item()
|
|
_, pred = output.max(1, keepdim=True)
|
|
_, label = target.max(dim=1)
|
|
correct += pred.eq(label.view_as(pred)).sum().item()
|
|
|
|
test_loss = test_loss
|
|
test_loss /= len(test_loader) # loss function already averages over batch size
|
|
print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
|
|
test_loss, correct, len(test_loader.dataset),
|
|
100. * correct / len(test_loader.dataset)))
|
|
|
|
|
|
test(net, test_dataloader, optimizer, criterion, device)
|