<jupyter_start><jupyter_text><jupyter_code>import Python
import TensorFlow<jupyter_output><empty_output><jupyter_text>Import some python libraries that we need<jupyter_code>%include "EnableIPythonDisplay.swift"
IPythonDisplay.shell.enable_matplotlib("inline")

let plt = Python.import("matplotlib.pyplot")
let np  = Python.import("numpy")
let subprocess = Python.import("subprocess")
let path = Python.import("os.path")<jupyter_output><empty_output><jupyter_text>Download cifar 10<jupyter_code>//https://github.com/tensorflow/swift-models/tree/master/CIFAR

let filepath = "./cifar-10-batches-py"
let isdir = Bool(path.isdir(filepath))!
if !isdir {
    print("Downloading CIFAR data...")
    let command = "wget -nv -O- https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz | tar xzf - -C ."
    subprocess.call(command, shell: true)
}<jupyter_output>Downloading CIFAR data...
2019-05-05 22:15:06 URL:https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz [170498071/170498071] -> "-" [1]<jupyter_text>Setup the dataset<jupyter_code>//https://github.com/tensorflow/swift-models/tree/master/CIFAR

var batchSize:Int = 64 

func loadCIFARFile(named name: String, in directory: String = ".") -> (Tensor<Float>, Tensor<Int32>) {
    let np = Python.import("numpy")
    let pickle = Python.import("pickle")
    let path = "\(directory)/cifar-10-batches-py/\(name)"
    let f = Python.open(path, "rb")
    let res = pickle.load(f, encoding: "bytes")

    let bytes = res[Python.bytes("data", encoding: "utf8")]
    let labels = res[Python.bytes("labels", encoding: "utf8")]

    let labelTensor = Tensor<Int64>(numpy: np.array(labels))!
    let images = Tensor<UInt8>(numpy: bytes)!
    let imageCount = images.shape[0]

    // reshape and transpose from the provided N(CHW) to TF default NHWC
    let imageTensor = Tensor<Float>(images
        .reshaped(to: [imageCount, 3, 32, 32])
        .transposed(withPermutations: [0, 2, 3, 1]))

    let mean = Tensor<Float>([0.485, 0.456, 0.406])
    let std  = Tensor<Float>([0.229, 0.224, 0.225])
    let imagesNormalized = ((imageTensor / 255.0) - mean) / std

    return (imagesNormalized, Tensor<Int32>(labelTensor))
}

/// helper functions 

// report accuracy of a batch 
func getAccuracy(y:Tensor<Int32>, logits:Tensor<Float>) -> Float{
  let out  = Tensor<Int32>(logits.argmax(squeezingAxis: 1) .== y).sum().scalarized()
  return Float(out) / Float(y.shape[0])
}

//round two decimal places 
func roundTwo(_ input:Float) -> Float{
  return (input*100).rounded()/100
}

//crop to a certain size 
func crop(_ tensor:Tensor<Float>, _ size:Int) -> Tensor<Float> {
  let i = Int.random(in: 0..<32-size)
  let j = Int.random(in: 0..<32-size)
  let N = Int(tensor.shape[0])
  
  return tensor[0..<N, i..<i+size, j..<j+size,0..<3]
}

//randomly augment a batch 
func augment(_ tensor:Tensor<Float>) -> Tensor<Float> {
  var out = tensor
  
  //maybe flip
  if Float.random(in:0...1) < 0.5{
    out = tensor.transposed(withPermutations: [0, 1, 2, 3])
  }
  //maybe crop and resize 
  if Float.random(in:0...1) < 0.5{
    let cropped = crop(tensor, 25)
    out = Raw.resizeArea(images:cropped , size:[32, 32] )
  }
  
  return out
}<jupyter_output><empty_output><jupyter_text>Create a dataset<jupyter_code>struct Element: TensorGroup {
    var x: Tensor<Float>
    var y: Tensor<Int32>
}

//cifar 10 training comes in 6 files we load/concatenate them [5000, 32, 32, 3]
let train_data = (1..<6).map { loadCIFARFile(named: "data_batch_\($0)") }

let trainX = Raw.concat(concatDim: Tensor<Int32>(0), train_data.map { $0.0})
let trainY = Raw.concat(concatDim: Tensor<Int32>(0), train_data.map { $0.1})

//load testing images size [1000, 32, 32, 3]
let (testX, testY) = loadCIFARFile(named: "test_batch")

//create a dataset for training and testing 
let trainDataset = Dataset<Element>(elements: Element(x: trainX, y:trainY))
let testDataset = Dataset<Element>(elements:  Element(x:testX, y:testY))<jupyter_output><empty_output><jupyter_text>Create the basic parts of the mode [convblocks + classifier]<jupyter_code>struct ConvBlock:Layer{

  typealias Input = Tensor<Float>
  typealias Output = Tensor<Float>
  
  var conv1: Conv2D<Float>
  var conv2: Conv2D<Float>
  var pool: MaxPool2D<Float>
  var norm: BatchNorm<Float>
  
  init(filterShape:(Int, Int))
  {
    self.conv1 = Conv2D<Float>(filterShape: (3, 3,filterShape.0, filterShape.1), 
                              strides: (1, 1), padding : .same, activation: relu)
    
    self.conv2 = Conv2D<Float>(filterShape: (3, 3,filterShape.1, filterShape.1), 
                              strides: (1, 1), padding : .same, activation: relu)
    
    self.norm = BatchNorm<Float>(featureCount: filterShape.1)
    self.pool = MaxPool2D<Float>(poolSize: (2, 2), strides: (2, 2))
  }
  
  @differentiable
  func call(_ input: Input) -> Output {
    return input.sequenced(through: conv1, conv2, norm, pool)
  }
}

struct Classifier:Layer{

  typealias Input = Tensor<Float>
  typealias Output = Tensor<Float>
  
  var dense1: Dense<Float>
  var dense2: Dense<Float>
  var dropout: Dropout<Float>
  
  init(input:Int, mid:Int)
  {
    self.dense1 = Dense<Float>(inputSize: input , outputSize: mid, activation: relu)
    self.dropout = Dropout<Float>(probability: 0.5)
    self.dense2 = Dense<Float>(inputSize: mid , outputSize: 10)
  }
  
  @differentiable
  func call(_ input: Input) -> Output {
    return input.sequenced(through: dense1, dropout, dense2)  
  }
}<jupyter_output><empty_output><jupyter_text>Create the overall model<jupyter_code>struct CNN: Layer {
    typealias Input = Tensor<Float>
    typealias Output = Tensor<Float>

    var conv1 = ConvBlock(filterShape:(3, 16))
    var conv2 = ConvBlock(filterShape:(16, 32))
    var conv3 = ConvBlock(filterShape:(32, 64))
    var conv4 = ConvBlock(filterShape:(64, 64))
  
    var dropout = Dropout<Float>(probability: 0.5)
  
    var flatten = Flatten<Float>()
    var classifier = Classifier(input:2*2*64, mid:128)
    
    @differentiable
    func call(_ input: Input) -> Output {
        let convolved = input.sequenced(through: conv1, conv2, conv3, conv4)
        return convolved.sequenced(through:dropout, flatten, classifier)
    }
}
var model = CNN()
let optimizer = Adam(for: model)

//warmup 
let tensor = Tensor<Float>(zeros: [1, 32, 32, 3])
print(model(tensor))<jupyter_output>[[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]]<jupyter_text>Training and reporting the results<jupyter_code>var trainLoss:Float = 0.0
var trainAcc :Float = 0.0
var testLoss:Float = 0.0
var testAcc:Float = 0.0 

var batchCount: Float = 0.0

for epoch in 0..<50{
  
  //evaluate metrics
  trainLoss = 0.0
  trainAcc  = 0.0
  batchCount = 0.0 
  
  let shuffled = trainDataset.shuffled(sampleCount:50000 , randomSeed: Int64(epoch))
    
  for batch in shuffled.batched(batchSize) {
  
    //get batches
    let X = augment(batch.x)
    let y = batch.y
    
    //calculate the loss and gradient
    let (loss, grads) = valueWithGradient(at: model) { model -> Tensor<Float> in
            let logits = model(X)
            return softmaxCrossEntropy(logits: logits, labels: y)
    }

    //make an optimizer step 
    optimizer.update(&model.allDifferentiableVariables, along: grads)    
    
    let logits = model(X) //this is slowing down ? 
    let acc = getAccuracy(y:y, logits:logits)
    
    trainLoss += Float(loss.scalarized())
    trainAcc  += acc
    batchCount += 1
  }  
 
  trainLoss /= batchCount
  trainAcc  /= batchCount
 
  //training
  testLoss = 0.0
  testAcc  = 0.0
  
  let logits = model(testX)
  let loss = softmaxCrossEntropy(logits: logits, labels: testY)
  let acc = getAccuracy(y:testY, logits:logits)

  testLoss += Float(loss.scalarized())
  testAcc  += acc
  print("epoch: \(epoch+1), train_loss: \(roundTwo(trainLoss)), test_loss: \(roundTwo(testLoss)), train_acc: \(roundTwo(trainAcc)), test_acc: \(roundTwo(testAcc))" )

}<jupyter_output>epoch: 1, train_loss: 1.64, test_loss: 1.3, train_acc: 0.41, test_acc: 0.53
epoch: 2, train_loss: 1.23, test_loss: 1.1, train_acc: 0.58, test_acc: 0.61
epoch: 3, train_loss: 1.04, test_loss: 0.94, train_acc: 0.66, test_acc: 0.67
epoch: 4, train_loss: 0.92, test_loss: 0.83, train_acc: 0.7, test_acc: 0.71
epoch: 5, train_loss: 0.84, test_loss: 0.87, train_acc: 0.73, test_acc: 0.7
epoch: 6, train_loss: 0.78, test_loss: 0.76, train_acc: 0.76, test_acc: 0.74
epoch: 7, train_loss: 0.74, test_loss: 0.76, train_acc: 0.77, test_acc: 0.74
epoch: 8, train_loss: 0.69, test_loss: 0.77, train_acc: 0.79, test_acc: 0.73
epoch: 9, train_loss: 0.66, test_loss: 0.72, train_acc: 0.8, test_acc: 0.75
epoch: 10, train_loss: 0.63, test_loss: 0.69, train_acc: 0.81, test_acc: 0.77
epoch: 11, train_loss: 0.61, test_loss: 0.67, train_acc: 0.82, test_acc: 0.77
epoch: 12, train_loss: 0.59, test_loss: 0.7, train_acc: 0.82, test_acc: 0.76
epoch: 13, train_loss: 0.58, test_loss: 0.71, train_acc: 0.83, test_acc: 0.77
e[...]
{'fraction_non_alphanumeric': 0.1488168666591903, 'fraction_numerical': 0.05203543792755411, 'mean_word_length': 3.8127361036157583, 'pattern_counts': {'":': 0, '<': 85, '<?xml version=': 0, '>': 85, 'https://': 4, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '8', 'dir': 'jupyter-structured-clean-dedup', 'id': '108208', 'license': 'MIT', 'n_tokens_mistral': 3663, 'n_tokens_neox': 3318, 'n_words': 781, 'path': 'Swift4TF_CIFAR10.ipynb', 'repo_name': 'Grace-TA/Colab-Notebook', 'size': '8917'}