print('CPU')
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
isAnomaly: torch.tensor
model = NeuralNetworkCalculator().to(device)
criterion = torch.nn.BCEWithLogitsLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
# scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.1)
for i in range(3, 3603):
tp = ls.load_image(DIRP + '/output%06d.jpg' % (i + 1)).to(device)
if os.path.isfile(DIRA + '/output%06d.jpg' % (i + 1)) == True:
ta = ls.load_image(DIRA + '/output%06d.jpg' % (i + 1)).to(device)
isAnomaly = torch.tensor([1, 0]).float().to(device)
else:
ta = ls.load_image(DIRNA + '/output%06d.jpg' % (i + 1)).to(device)
isAnomaly = torch.tensor([0, 1]).float().to(device)
t = torch.cat((tp, ta))
t = t.view(1, 1, 2, 480, 640)
optimizer.zero_grad()
output = model(t)
loss = criterion(output, isAnomaly)
loss.backward()
optimizer.step()
# Train the model
for epoch in range(num_epochs):
datasamples = len([
name for name in os.listdir(DIRT)
if os.path.isdir(os.path.join(DIRT, name))
])
for folder in range(datasamples):
DIRRT = DIRT + '/' + str(folder + 1)
total_step = len([
name for name in os.listdir(DIRRT)
if os.path.isfile(os.path.join(DIRRT, name))
]) - 6
for file in range(total_step):
imageinN1 = loadsave.load_image(DIRRT + '/output%06d.jpg' %
(file + 1)).to(device)
imageinN2 = loadsave.load_image(DIRRT + '/output%06d.jpg' %
(file + 2)).to(device)
"""imageinN3 = loadsave.load_image(DIRRT+'/output%06d.jpg' %(file+3)).to(device)
imageinN4 = loadsave.load_image(DIRRT+'/output%06d.jpg' %(file+4)).to(device)
imageinN5 = loadsave.load_image(DIRRT+'/output%06d.jpg' %(file+5)).to(device)"""
imageout = loadsave.load_image(DIRRT + '/output%06d.jpg' %
(file + 3)).to(device)
frame_total = frame_total + 1
outputs = model(imageinN1,
imageinN2) #, imageinN3, imageinN4, imageinN5
loss = criterion(outputs, imageout)
# Backward and optimize