Skip to content

Instantly share code, notes, and snippets.

@seanbenhur
Last active September 2, 2020 07:13
Show Gist options
  • Save seanbenhur/749aede8c1f2d5c1bb0862fee0718162 to your computer and use it in GitHub Desktop.
Save seanbenhur/749aede8c1f2d5c1bb0862fee0718162 to your computer and use it in GitHub Desktop.
# Declare Siamese Network
net = SiameseNetwork().cuda()
# Decalre Loss Function
criterion = ContrastiveLoss()
# Declare Optimizer
optimizer = th.optim.Adam(net.parameters(), lr=1e-3, weight_decay=0.0005)
#train the model
def train():
loss=[]
counter=[]
iteration_number = 0
for epoch in range(1,config.epochs):
for i, data in enumerate(train_dataloader,0):
img0, img1 , label = data
img0, img1 , label = img0.cuda(), img1.cuda() , label.cuda()
optimizer.zero_grad()
output1,output2 = net(img0,img1)
loss_contrastive = criterion(output1,output2,label)
loss_contrastive.backward()
optimizer.step()
print("Epoch {}\n Current loss {}\n".format(epoch,loss_contrastive.item()))
iteration_number += 10
counter.append(iteration_number)
loss.append(loss_contrastive.item())
show_plot(counter, loss)
return net
#set the device to cuda
device = torch.device('cuda' if th.cuda.is_available() else 'cpu')
model = train()
torch.save(model.state_dict(), "model.pt")
print("Model Saved Successfully")
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment