som = SOM(input_size=28 * 28 * 1, out_size=(row, col))
if os.path.exists('%s/som.pth' % MODEL_DIR):
som.load_state_dict(torch.load('%s/som.pth' % MODEL_DIR))
print('Model Loaded!')
else:
print('Create Model!')
som = som.to(device)
if train == True:
losses = list()
for epoch in range(total_epoch):
running_loss = 0
start_time = time.time()
for idx, (X, Y) in enumerate(train_loader):
X = X.view(-1, 28 * 28 * 1).to(device) # flatten
loss = som.self_organizing(X, epoch, total_epoch) # train som
running_loss += loss
losses.append(running_loss)
print('epoch = %d, loss = %.2f, time = %.2fs' %
(epoch + 1, running_loss, time.time() - start_time))
if epoch % 5 == 0:
# save
som.save_result('%s/som_epoch_%d.png' % (RES_DIR, epoch),
(1, 28, 28))
torch.save(som.state_dict(), '%s/som.pth' % MODEL_DIR)
torch.save(som.state_dict(), '%s/som.pth' % MODEL_DIR)
plt.title('SOM loss')
plt.plot(losses)
colors = np.append(colors, np.array([[0, 0, 0]]), axis=0)
colors = np.append(colors, np.array([[1, 1, 1]]), axis=0)
for i in range(10):
colors = np.append(colors, np.array([[0, 0, random.random()]]), axis=0)
colors = np.append(colors, np.array([[0, random.random(), 0]]), axis=0)
colors = np.append(colors, np.array([[random.random(), 0, 0]]), axis=0)
colors = np.append(colors, np.array([[1, 1, random.random()]]), axis=0)
colors = np.append(colors, np.array([[1, random.random(), 1]]), axis=0)
colors = np.append(colors, np.array([[random.random(), 1, 1]]), axis=0)
colors = np.append(colors, np.array([[0, random.random(), random.random()]]), axis=0)
colors = np.append(colors, np.array([[random.random(), random.random(), 0]]), axis=0)
colors = np.append(colors, np.array([[1, random.random(), random.random()]]), axis=0)
colors = np.append(colors, np.array([[random.random(), random.random(), 1]]), axis=0)
colors = np.append(colors, np.array([[random.random(), random.random(), random.random()]]), axis=0)
data = torch.Tensor(colors)
row = 40
col = 40
total_epoch = 1000
som = SOM(3, (row, col))
for iter_no in range(total_epoch):
som.self_organizing(data, iter_no, total_epoch)
weight = som.weight.reshape(3, row, col).numpy()
weight = np.transpose(weight, (1, 2, 0,))
plt.title('Color SOM')
plt.imshow(weight)
plt.show()
