---
# System prepended metadata

title: Untitled

---

import torch
import torchvision
import torchvision.transforms as transforms
import torch.nn as nn
import torch.optim as optim
import datetime


import math






if __name__=='__main__':
	assert torch.cuda.is_available()
	transform=transforms.Compose([
		transforms.Resize(256),
		transforms.CenterCrop(224),
		transforms.ToTensor(),
		transforms.Normalize(mean=(0.485,0.456,0.406),std=(0.229,0.224,0.225)),
	])
	train_set=torchvision.datasets.CIFAR10(root='./data',train=True,transform=transform)
	test_set=torchvision.datasets.CIFAR10(root='./data',train=False,transform=transform)
	model=torch.hub.load('pytorch/vision:v0.10.0','resnet34',pretrained=False)
	model.fc=nn.Linear(model.fc.in_features,10)
	model=model.to('cuda')
    
	start_datetime=datetime.datetime.now()
	step_lr_decay=0
	lr_now=0.1
	sqrtannealing=True
	start_batchsize=256
    
	criterion=nn.CrossEntropyLoss()
	optimizer=optim.SGD(model.parameters(),lr=lr_now,weight_decay=1e-3,momentum=0.9)
	scheduler=optim.lr_scheduler.ReduceLROnPlateau(optimizer,mode='min',factor=0.125,patience=1,verbose=True)
	train_log=[]
	test_log=[]
	for i in range(300):
		batchsize=start_batchsize//2**step_lr_decay if sqrtannealing else start_batchsize
		batchsize=max(batchsize,1)
		train_loader=torch.utils.data.DataLoader(train_set,batch_size=batchsize,shuffle=True)
		test_loader=torch.utils.data.DataLoader(test_set,batch_size=start_batchsize*2,shuffle=True)
		model.train()
		loss_sum,cnt,n=0,0,len(train_loader.dataset)
		for data,target in train_loader:
			data,target=data.to('cuda'),target.to('cuda')
			optimizer.zero_grad()
			output=model(data)
			loss=criterion(output,target)
			loss.backward()
			nn.utils.clip_grad_norm_(model.parameters(), 10.0)
			optimizer.step()
			loss_sum+=loss.item()*len(data)
			cnt+=len(data)
			print('[{}]{:3}: {:.6f}  {:3}/{:3}({:4.0%})'.format(
				str(datetime.datetime.now()-start_datetime).split('.')[0],
				i+1,loss_sum/cnt,cnt,n,cnt/n
			),end='\r',flush=True)
		print('[{}]{:3}: train loss = {:.12f}'.format(
			str(datetime.datetime.now()-start_datetime).split('.')[0],
			i+1,loss_sum/cnt
		),flush=True)
		train_log.append(loss_sum/cnt)
		model.eval()
		test_loss,cnt,n=0,0,len(test_loader.dataset)
		with torch.no_grad():
			for data,target in test_loader:
				data,target=data.to('cuda'),target.to('cuda')
				output=model(data)
				loss=criterion(output,target)
				test_loss+=loss.item()*len(data)
				cnt+=len(data)
				print('[{}]{:3}: {:.6f}  {:3}/{:3}({:4.0%})'.format(
					str(datetime.datetime.now()-start_datetime).split('.')[0],
					i+1,test_loss/cnt,cnt,n,cnt/n
				),end='\r',flush=True)
		print('[{}]{:3}:  test loss = {:.12f}'.format(
			str(datetime.datetime.now()-start_datetime).split('.')[0],
			i+1,test_loss/cnt
		),flush=True)
		scheduler.step(test_loss)
        test_log.append(test_loss/cnt)
		assert len(scheduler._last_lr)==1
		if scheduler._last_lr[0]<lr_now :
			print('[{}]{:3}: lr {:.12f} -> {:.12f}'.format(
				str(datetime.datetime.now()-start_datetime).split('.')[0],
				i+1,lr_now,scheduler._last_lr[0]
			),flush=True)
			lr_now=scheduler._last_lr[0]
			step_lr_decay+=1
			#print("ne　wbatchsize=",batchsize)
            if sqrtannealing:
                assert len(optimizer.param_groups)==1
                m_old=optimizer.param_groups[0]['momentum']
                x=2
                #m_new=(m_old-1)/math.sqrt(x)+1
                m_new=max(1-math.sqrt(x)*(1-m_old),0.1)
                optimizer.param_groups[0]['momentum']=m_new
                print(optimizer.param_groups[0]['momentum'])

		if (i+1)%10==0:
			model.eval()
			print(step_lr_decay)
			correct,cnt=[0]*10,[0]*10
			#print(train_loss)
			with torch.no_grad():
				for data,target in test_loader:
					data,target=data.to('cuda'),target.to('cuda')
					output=model(data)
					_,pred = torch.max(output,1)
					pred=pred==target
					for j in range(10):
						cnt[j]+=(target==j).sum()
						correct[j]+=pred[target==j].sum()
			print('[{}]{:3}:'.format(
				str(datetime.datetime.now()-start_datetime).split('.')[0],i+1
		　	),flush=True)
			for j in range(10):
				print('{:10} : {:5}/{:5}({:4.0%})'.format(
					train_set.classes[j],correct[j],cnt[j],correct[j]/cnt[j]
				),flush=True)
			print('{:10} : {:5}/{:5}({:4.0%})'.format(
				'all',sum(correct),sum(cnt),sum(correct)/sum(cnt)
			),flush=True)
	print(train_log)
	print(test_log)