def main(config):
# Use GPU!
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print("Device: " + str(device))
# Simple transform
transform = transforms.Compose([transforms.ToTensor()]) #transforms.Normalize((0.1307,), (0.3081,)
# Get MNIST Datasets
save = config.data_loc #'./data/MNIST' #'/home/kartik/git/ADLG/Assignment2/data/MNIST'
trainloader, testloader,_ = LoadMNIST(save, transforms_=transform,batch_size=config.batch_size)
epochs = config.epochs
print(config.epochs)
model = AutoEncoder()
# opt = optim.Adam(model.parameters(), lr=0.0005, betas=(0.9, 0.95))
opt = optim.SGD(model.parameters(), lr=config.lr)
# Loss function
# cost = config.loss
cost = eval(config.loss + "()") # nn.MSELoss()
loss_name = eval(config.loss + '.__name__')
# cost = nn.BCELoss() #reduce=False
decode_out = config.save_images + loss_name + '/' #'./save/Decoder_Images/'
# Create a trainer
trainer = Trainer(model, opt, cost, name="Default Autoencoder",
device=device, decode_out=decode_out)
# Run training
trainer.Train(trainloader, epochs)
#################################
# Optimizer
opt = optim.Adam(model.parameters(), lr=0.0005, betas=(0.9, 0.95))
# Loss function
cost = nn.CrossEntropyLoss()
# Create a trainer
trainer = Trainer(model, opt, cost, name="Default CNN", device=device)
# Add test accuracy plotting
plotter.new_line("Default CNN")
trainer.SetEpochCallback(plotter.EpochCallback)
# Run training
trainer.Train(trainloader, epochs, testloader=testloader)
#################################
# Create the assignment Resnet (part a)
#################################
def GetDefaultResNet():
resnet = ResNet(in_features= [32, 32, 3],
num_class=10,
feature_channel_list = [32, 64, 128],
batch_norm= True,
)
# Create MLP
# Calculate the input shape
s = resnet.GetCurShape()
# Send model to GPU
model.to(device)
# Specify optimizer
opt = optim.Adam(model.parameters(), lr=0.0005, betas=(0.9, 0.95))
# Specify loss function
cost = nn.CrossEntropyLoss()
# Train the model
trainer = Trainer(device=device, name="Basic_CNN")
epochs = 5
trainer.Train(model,
trainloader,
testloader,
cost=cost,
opt=opt,
epochs=epochs)
# Load best saved model for inference
model_loaded = GetCNN()
# Specify location of saved model
PATH = "./save/Basic_CNN-best-model/model.pt"
checkpoint = torch.load(PATH)
# load the saved model
model_loaded.load_state_dict(checkpoint['state_dict'])
# intialization for hooks and storing activation of ReLU layers
activation = {}