func main() {
training := mnist.LoadTrainingDataSet()
samples := make([][]bool, len(training.Samples))
for i, sample := range training.Samples {
samples[i] = make([]bool, len(sample.Intensities))
for j, x := range sample.Intensities {
if x > 0.5 {
samples[i][j] = true
}
}
}
layers := buildLayers()
trainer := rbm.Trainer{
GibbsSteps: GibbsSteps,
StepSize: BigStepSize,
Epochs: BigEpochs,
BatchSize: runtime.GOMAXPROCS(0),
}
log.Println("Training...")
trainer.TrainDeep(layers, samples[:1000])
trainer.StepSize = SmallStepSize
trainer.Epochs = SmallEpochs
trainer.TrainDeep(layers, samples[:1000])
log.Println("Generating outputs...")
testingSamples := mnist.LoadTestingDataSet()
mnist.SaveReconstructionGrid("output.png", func(img []float64) []float64 {
return reconstruct(layers, img)
}, testingSamples, ReconstructionGridSize, ReconstructionGridSize)
}
