# สร้าง class ขึ้นมาแล้วเรียกใช้
from torch import nn
loss_func = nn.CrossEntropyLoss()
loss_func(acts, targ)

# เรียกใช้ functional API
import torch.nn.functional as F
F.cross_entropy(acts, targ)

#ถ้าอยากได้ตัวเลขก่อน aggregation (ส่วนใหญ๋ aggregate ด้วย mean)
nn.CrossEntropyLoss(reduction='none')(acts, targ)

#สร้างตัวแปร X, y สำหรับเทรนโมเดล 
#ต่อ tensor สองอันเข้าด้วยกัน (default ด้วยมิติ 0)
train_x = torch.cat([stacked_threes, stacked_sevens])\
	.view(-1, 28*28) #เปลี่ยนมิติเป็น (-1, 28*28); -1 คือให้มีมิติเท่าเดิม (จำนวนรูป)

#ให้เลข 3 เป็น 1 ส่วนเลข 7 เป็น 0
train_y = tensor([1]*len(threes) + [0]*len(sevens))\
	.unsqueeze(1) #เพิ่ม dimension ให้กับ tensor; ในกรณีนี้คือเจาะจงว่าให้เพิ่มใน dimension ที่ 1 คือ (# examples,) -> (# examples, 1)

#สร้าง dataset object ไว้เพื่อเก็บคู่ (X,y) ของ training set
dset = list(zip(train_x,train_y)) #zip (X,y) เข้าด้วยกันแล้วเปลี่ยนเป็น list
x,y = dset[0] #เลือก (X,y) ที่ index 0

#ทำแบบเดียวกันกับ validation set
valid_x = torch.cat([valid_3_tens, valid_7_tens]).view(-1, 28*28)
valid_y = tensor([1]*len(valid_3_tens) + [0]*len(valid_7_tens)).unsqueeze(1)
valid_dset = list(zip(valid_x,valid_y))

#ตั้งค่า weights และ bias เริ่มต้นด้วยการสุ่มจาก standard normal distribution
def init_params(size, std=1.0): 
	return (torch.randn(size)*std).requires_grad_() #requires_grad_() เพื่อให้ Pytorch เก็บ gradient สำหรับ backpropagation

#สมการของเราคือ y = weights * X + bias
weights = init_params((28*28,1))
bias = init_params(1)

#prediction สำหรับ example #0
#ผลรวมของ (# examples, 784) * (1, 784) บวกด้วย (1,)
#weights.T มีมิติ (1, 784) จะถูก broadcast ไปคูณสำหรับทุก example 
#ดูวิธีทำได้จาก http://matrixmultiplication.xyz/
(train_x[0] * weights.T).sum() + bias 

#สำหรับ Pytorch @ คือการทำ matrix multiplication ตามวิธีข้างต้น
def linear1(xb): return xb@weights + bias
preds = linear1(train_x) #dimension คือ (# examples, 1)

corrects = (preds>0.0).float() == train_y #ถ้า prediction > 0 ให้ทายว่า 1 (class ของ 7) แล้วดูว่าตรงกับ ground truth (y) แค่ไหน (True/False)

#จากโมเดลที่ไม่ได้เทรนเลยเราจะเห็นได้ว่าตรงประมาณครึ่งๆ (เพราะเรามีเลข 3 และ 7 พอๆกันถ้าเดามั่วๆมีโอกาสถูกครึ่งๆ)
corrects.float().mean().item() #item() เพื่อเอาค่าที่ถูกเก็บไว้ใน tensor ไม่ใช่ตัว tensor ทั้งหมด

#หนึ่งในเหตุผลที่เราไม่ใช้ accuracy เป็น loss function คือ accuracy ไม่ไวต่อการเปลี่ยนแปลงของ weights/bias พอ เช่น ในกรณีนี้หากเราคูณ weights ด้วย 1.0001 เราก็จะยังได้ accuracy เท่าเดิม (เพราะ threshold ที่จะเปลี่ยนคำทำนายของเราคือ prediction > 0.0) ทำให้ยากที่โมเดลจะเรียนรู้
weights[0] *= 1.0001
preds = linear1(train_x)
((preds>0.0).float() == train_y).float().mean().item()

#ถ้า ground truth คือ 1 ให้ค่า 1-predictions ถ้าเป็น 0 ให้ค่า predictions
#prediction ยิ่งใกล้ ground truth ที่ถูกเท่าไหร่ ค่า loss ยิ่งน้อยลงเท่านั้น
#ค่า loss เปลี่ยนทุกครั้งที่ prediction เปลี่ยน ไม่เหมือนกับ accuracy ด้านบน
def mnist_loss(predictions, targets):
    return torch.where(targets==1, 1-predictions, predictions).mean()

#ปัญหาคือถ้า predictions ไม่อยู่ระหว่าง 0 กับ 1 ค่า loss อาจจะไม่สมเหตุสมผล
#เช่น predictions = 100, ground truth = 1, loss = 1-100 = -99
#เราจึงใช้ activation function เช่น sigmoid
def sigmoid(x): return 1/(1+torch.exp(-x))

#เราใส่ sigmoid function ไปใน loss เลยก็ได้

def mnist_loss(predictions, targets):
    predictions = predictions.sigmoid()
    return torch.where(targets==1, 1-predictions, predictions).mean()

coll = range(15) #ข้อมูลเลขจาก 0-14
#dl คือ iterator
dl = DataLoader(coll, batch_size=5, shuffle=True) #โหลดข้อมูลทีละ 5 ตัวอย่างต่อ 1 mini-batch; สับข้อมูลมั่วๆด้วย
list(dl) #เปลี่ยน iterator เป็น list ทำให้แสดงข้อมูลทั้งหมด 3 mini-batches; mini-batch ละ 5 ตัวอย่าง

#dataloader สำหรับ tuples
ds = L(enumerate(string.ascii_lowercase)) #[(0,'a'), (1,'b'),...]
dl = DataLoader(ds, batch_size=3, shuffle=False)
list(dl) #[((0,1,2),('a','b','c')), ...]

def calc_grad(xb, yb, model):
    preds = model(xb)
    loss = mnist_loss(preds, yb)
    loss.backward() #คำนวณ gradients จาก operations ด้านบน

def train_epoch(model, lr, params):
    for xb,yb in dl:
        calc_grad(xb, yb, model) #คำนวณ gradients
        #สำหรับ parameter ทุกตัว
        for p in params:
        	#ลบ parameter ด้วยค่า gradeint คูณด้วย learning rate
        	#ใช้ p.data เพื่อไม่ให้ Pytorch คำนวณ gradient จากการลบนี้
            p.data -= p.grad*lr 
            #จำเป็นต้องเซ็ต gradients จากทุก tensor เป็น 0 ไม่งั้นเวลาเรียก .backward() จะเอา gradients ใหม่มาบวกกับอันเก่า
            p.grad.zero_() 

def batch_accuracy(xb, yb):
    preds = xb.sigmoid()
    correct = (preds>0.5) == yb
    return correct.float().mean()

def validate_epoch(model):
    accs = [batch_accuracy(model(xb), yb) for xb,yb in valid_dl]
    return round(torch.stack(accs).mean().item(), 4) #คำนวณค่าเฉลี่ยของ acc จากทุก mini-batch

lr = 1. #learning rate
params = weights,bias #parameter ของเรา ในกรณีนี้มีแค่ weights และ bias ที่เอามาทำ y=xb@weights + bias

for i in range(20): #เทรนไป 20 epochs (เทรนด้วยข้อมูลทั้งหมด 20 รอบ)
    train_epoch(linear1, lr, params)
    print(validate_epoch(linear1), end=' ')

linear_model = nn.Linear(28*28,1) #สร้างโมเดลที่ทำ y=xb@weights + bias ง่ายๆ

#นิยาม weights และ bias ด้วย .parameters()
w,b = linear_model.parameters()

#สร้างคลาส optimizer
class BasicOptim:
    def __init__(self,params,lr): 
    	#กำหนด parameters และ learning rate เบื้องต้น
    	self.params = list(params)
    	self.lr = lr

    def step(self, *args, **kwargs):
    	#ทำ SGD
        for p in self.params: p.data -= p.grad.data * self.lr

    def zero_grad(self, *args, **kwargs):
    	#รีเซต gradients ให้เป็น 0
        for p in self.params: p.grad = None

#สร้าง optimizer
opt = BasicOptim(linear_model.parameters(), lr)

#train ด้วย optimizer ที่เราสร้าง
def train_epoch(model):
    for xb,yb in dl:
        calc_grad(xb, yb, model)
        opt.step()
        opt.zero_grad()

#เทรนโมเดลเป็นจำนวน X epochs
def train_model(model, epochs):
    for i in range(epochs):
        train_epoch(model)
        print(validate_epoch(model), end=' ')

#initialize weights/bias
linear_model = nn.Linear(28*28,1)
#สร้าง optimizer
opt = SGD(linear_model.parameters(), lr) #Pytorch/fastai มีคลาส SGD ที่เป็นคลาส optimizer เหมือนกันที่เราทำข้างบนอยู่แล้ว
#เทรนโมเดล 20 epochs
train_model(linear_model, 20)
#สร้าง dataloaders
dls = DataLoaders(dl, valid_dl)
#สร้าง learner object ที่ทำการรวบทุกอย่างเข้าด้วยกันเพื่อความง่าย
learn = Learner(dls, #dataloaders
	nn.Linear(28*28,1), #model
	opt_func=SGD, #optimizer
	loss_func=mnist_loss, #loss function
	metrics=batch_accuracy #metric
	)
#เทรนโมเดลจาก learner object
learn.fit(10, lr=lr)

def simple_net(xb): 
    res = xb@w1 + b1
    res = res.max(tensor(0.0)) #ReLU หาค่า max ระหว่างค่าที่อยู่ในแต่ละ element ของ res กับ 0
    res = res@w2 + b2
    return res

#ทำแบบเดียวกันด้วย Pytorch ก็ได้
simple_net = nn.Sequential(
    nn.Linear(28*28,30),
    nn.ReLU(),
    nn.Linear(30,1)
)

#สร้าง learner แล้วทำเหมือนเดิม
learn = Learner(dls, simple_net, opt_func=SGD,
                loss_func=mnist_loss, metrics=batch_accuracy)

#เทรน
learn.fit(40, 0.1)

#โชว์ผล batch_accuracy ตลอด 40 epochs
plt.plot(L(learn.recorder.values).itemgot(2)); #get item at index 2 from each row

m = learn.model
w,b = m[0].parameters() #parameters ของ layer แรก
show_image(w[0].view(28,28))

pets = DataBlock(blocks = (ImageBlock, CategoryBlock),
                 get_items=get_image_files, 
                 splitter=RandomSplitter(seed=42),
                 get_y=using_attr(RegexLabeller(r'(.+)_\d+.jpg$'), 'name'),
                 item_tfms=Resize(460), #crop and resize randomly first
                 batch_tfms=aug_transforms(size=224, min_scale=0.75)) #augmentation transforms
dls = pets.dataloaders(path/"images")

#เช็คว่า data augmentation สมเหตุสมผล
dls.show_batch(nrows=1, ncols=3)
#เช็คว่า data augmentation สมเหตุสมผลแบบดูจากรูปเดียว
dls.show_batch(nrows=1, ncols=3, unique=True)

pets1 = DataBlock(blocks = (ImageBlock, CategoryBlock),
                 get_items=get_image_files, 
                 splitter=RandomSplitter(seed=42),
                 get_y=using_attr(RegexLabeller(r'(.+)_\d+.jpg$'), 'name'))

#ตัวอย่างที่ไม่ resize ด้วย item_tfms ก่อน ทำให้ DataBlock เอารูปมาต่อกันไม่ได้
pets1.summary(path/"images")

#softmax activation ให้ "probability" ของแต่ละคลาสที่รวมกันแล้วได้ 1.
def softmax(x): return exp(x) / exp(x).sum(dim=1, keepdim=True)
sm_acts = torch.softmax(acts, dim=1) #apply softmax โดยให้แต่ละแถวรวมกันได้ 1

#เลือก activation ใน column ของค่า target; ตัวที่ 1 เลือก column 0, ตัวที่ 2 เลือก column 1, ...
targ = tensor([0,1,0,1,1,0])
idx = range(6)
sm_acts[idx, targ]

#Pytorch-sytle negative log likelihood loss
-sm_acts[idx, targ]

#Pytorch version; จะเห็นได้ว่าถึงมี log ในชื่อแต่มันเป็นแค่ sum of -softmax(x) ไม่มี log
#ทั้งนี้เป็นเพราะ Pytorch คาดหวังให้เราใช้ torch.log_softmax ที่จะทำ log(softmax(x)) ก่อน
F.nll_loss(sm_acts, targ, reduction='none')

#CrossEntropyLoss ใน Pytorch จะทำการหา sum of -log(softmax(x))
loss_func = nn.CrossEntropyLoss()