def main():
args = parse_args()
print("device is available : ", torch.cuda.is_available())
print("Start.....")
checkpoint = os.path.join('img', args.checkpoint)
model = load_model(args, checkpoint)
video = os.path.join("input2.avi")
video_reading(model, video)
print("Finish")
model = model.cuda()
model.train()
# Train!
for epoch in range(0, args.epochs):
bar = tqdm(train_data_loader)
for inputs, labels in bar:
inputs = inputs.cuda()
# labels = labels.cuda()
model.forward(inputs)
model.backward()
loss_G, loss_D = model.getLoss()
bar.set_description("Loss_G: " + str(loss_G) + " loss_D: " +
str(loss_D))
bar.refresh()
torch.save({
'epoch': epoch + 1,
'state_dict': model.state_dict(),
}, os.path.join('save', 'Ganormaly_epoch-' + str(epoch) + '.pth.tar'))
print("Save model at {}\n".format(
os.path.join('save',
'Ganormaly_epoch-' + str(epoch) + '.pth.tar')))
print(len(train_data_loader), loss_G / len(train_data_loader))
if __name__ == '__main__':
args = parse_args(phase='train')
train(args)
# sunnerTransforms.Transpose(),
sunnerTransforms.Normalize(mean=[0.5, 0.5, 0.5],
std=[0.5, 0.5, 0.5]),
])),
batch_size=args.batch_size,
shuffle=True,
num_workers=2)
# Create the model
model = GANomaly2D(r=args.r, device=args.device)
model.IO(args.resume, direction='load')
# Demo!
bar = tqdm(loader)
#for i, (abnormal_img) in enumerate(bar):
# # model.forward(abnormal_img)
# # model.backward()
# if i % args.record_iter == 0:
model.eval()
with torch.no_grad():
for (img, ) in bar:
z, z_ = model.forward(img)
img, img_ = model.getImg()
visualizeAnomalyImage(img, img_, z, z_)
model.demo()
if __name__ == '__main__':
args = parse_args(phase='demo')
demo(args)
"""
return self.x, self.x_
if __name__ == '__main__':
import sys
import os
from torch.utils.data import DataLoader
# __file__获取执行文件相对路径，整行为取上一级的上一级目录
BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.append(BASE_DIR)
print(BASE_DIR)
from parse.parse import parse_args
from lib.dataloader import VideoDataset
args = parse_args()
model = GANomaly3D(args=args)
model = model.cuda()
train_data = VideoDataset(args)
train_loader = DataLoader(train_data,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers)
for i, sample in enumerate(train_loader):
inputs = sample[0].cuda()
labels = sample[1]
print(inputs.size())
print(labels)