input_chiba <- layer_input(shape = c(28,28,1))

brach_a = input_chiba %>% 
  layer_conv_2d(filter = 32, kernel_size = 1,
                activation = "relu", strides = 2)

brach_b = input_chiba %>% 
  layer_conv_2d(filter = 32, kernel_size = 1,
                activation = "relu") %>% 
  layer_conv_2d(filter = 32, kernel_size = 2,
                activation = "relu", stride = 2)

brach_c = input_chiba %>% 
  layer_average_pooling_2d(pool_size = 2, stride = 2) %>% 
  layer_conv_2d(filter = 32, kernel_size = 1,
                activation = "relu")

brach_d = input_chiba %>% 
  layer_conv_2d(filter = 32, kernel_size = 1,
                activation = "relu") %>% 
  layer_conv_2d(filter = 32, kernel_size = 1,
                activation = "relu") %>% 
  layer_conv_2d(filter = 32, kernel_size = 2,
                activation = "relu", stride = 2)

concat = layer_concatenate(list(brach_a, brach_b, brach_c, brach_d))

output_chiba = concat %>% layer_flatten() %>%
  layer_dense(units = 64, activation = "relu") %>%
  layer_dense(units = 10, activation = "softmax")

chiba_model <- keras_model(input_chiba, output_chiba)

mnist <- dataset_mnist()
c(c(train_images, train_labels),c(test_images,test_labels)) %<-% mnist

train_images <- array_reshape(train_images,c(60000,28,28,1))
test_images <- array_reshape(test_images,c(10000,28,28,1))

train_images = train_images/255
test_images = test_images/255

train_labels = to_categorical(train_labels)
test_labels = to_categorical(test_labels)

chiba_model %>% compile(
  optimizer = "rmsprop",
  loss = "categorical_crossentropy",
  metrics = c("accuracy")
)

history <- chiba_model %>% fit(
  train_images,
  train_labels,
  epochs = 15,
  batch_size =64,
  validation_data = list(test_images,test_labels)
)
plot(history)