EPOCH = 20
BATCH_SIZE = 128
LR = 0.005 # learning rate
torch.cuda.empty_cache()
data_transforms = torchvision.transforms.Compose([
torchvision.transforms.RandomResizedCrop(344),
torchvision.transforms.RandomHorizontalFlip(),
torchvision.transforms.ToTensor()])
path1 = ‘drive/My Drive/Colab/image/test/’
train_data = torchvision.datasets.ImageFolder(path1, transform=data_transforms)
class AutoEncoder(nn.Module):
def init(self):
super(AutoEncoder, self).init()
self.encoder = nn.Sequential(
nn.Linear(3*344*344, 128),
nn.Tanh(), # 激活
nn.Linear(128, 64),
nn.Tanh(),
nn.Linear(64, 12),
nn.Tanh(),
nn.Linear(12, 3), # compress to 3 features which can be visualized in plt
)
self.decoder = nn.Sequential(
nn.Linear(3, 12),
nn.Tanh(),
nn.Linear(12, 64),
nn.Tanh(),
nn.Linear(64, 128),
nn.Tanh(),
nn.Linear(128, 3*344*344),
nn.Sigmoid(), # compress to a range (0, 1)
)
def forward(self, x):
x = x.view(x.size(0), -1)
encoded = self.encoder(x)
decoded = self.decoder(encoded)
return encoded, decoded
autoencoder = AutoEncoder()
optimizer = torch.optim.Adam(autoencoder.parameters(), lr=LR)
loss_func = nn.MSELoss()
initialize figure
f, a = plt.subplots(2, N_TEST_IMG, figsize=(5, 2))
plt.ion() # continuously plot
original data (first row) for viewing
view_data = train_data.train_data[:N_TEST_IMG].view(-1, 344*344).type(torch.FloatTensor)/255.
for i in range(N_TEST_IMG):
a[0][i].imshow(np.reshape(view_data.data.numpy()[i], (344, 344)), cmap=‘rainbow’); a[0][i].set_xticks(()); a[0][i].set_yticks(())
for epoch in range(EPOCH):
for step, (x, b_label) in enumerate(train_loader):
b_x = x.view(-1, 3344344) # batch x, shape (batch, 2828)
b_y = x.view(-1, 3344344) # batch y, shape (batch, 2828)
encoded, decoded = autoencoder(b_x)
loss = loss_func(decoded, b_y) # mean square error
optimizer.zero_grad() # clear gradients for this training step
loss.backward() # backpropagation, compute gradients
optimizer.step() # apply gradients
Above is the training process
with torch.no_grad():
for img, label in train_loader :
fig = plt.figure()
print(‘img’, img.shape)
imggg = np.transpose(img[0],(1,2,0))
print(‘imggg’, imggg.shape)
ax1 = fig.add_subplot(121)
ax1.imshow(imggg)
if torch.cuda.is_available():
img = Variable(img.to())
else:
img = Variable(img)
encoded, decoded = autoencoder(img)
print(decoded.shape)
print(decoded)
decodeddd = np.transpose(decoded.cpu()[0],(1,2,0))
print(decodeddd.shape)
print(decodeddd)
ax2 = fig.add_subplot(122)
ax2.imshow(decodeddd)
this is the code that i used to plot my decoded picture, but it returns the error:
TypeError Traceback (most recent call last)
/usr/local/lib/python3.6/dist-packages/numpy/core/fromnumeric.py in _wrapfunc(obj, method, *args, **kwds)
55 try:
—> 56 return getattr(obj, method)(*args, **kwds)
57
TypeError: transpose(): argument ‘dim0’ (position 1) must be int, not tuple
During handling of the above exception, another exception occurred:
ValueError Traceback (most recent call last)
3 frames
/usr/local/lib/python3.6/dist-packages/numpy/core/fromnumeric.py in _wrapit(obj, method, *args, **kwds)
44 except AttributeError:
45 wrap = None
—> 46 result = getattr(asarray(obj), method)(*args, **kwds)
47 if wrap:
48 if not isinstance(result, mu.ndarray):
ValueError: axes don’t match array
How to solve the problem, or what should i do to show my decoded pic?