train_loader = DataLoader(train_set,
batch_size=batch_size,
shuffle=True,
num_workers=16,
pin_memory=True)
val_loader = DataLoader(val_set,
batch_size=batch_size,
shuffle=True,
num_workers=16,
pin_memory=True)
# Create model
# model = CNN3D(sample_size=sample_size, sample_duration=sample_duration, drop_p=drop_p,
# hidden1=hidden1, hidden2=hidden2, num_classes=num_classes).to(device)
model = resnet18(pretrained=True,
progress=True,
sample_size=sample_size,
sample_duration=sample_duration,
attention=attention,
num_classes=num_classes).to(device)
# model = r2plus1d_18(pretrained=True, num_classes=num_classes).to(device)
# Run the model parallelly
if torch.cuda.device_count() > 1:
logger.info("Using {} GPUs".format(torch.cuda.device_count()))
model = nn.DataParallel(model)
# Create loss criterion & optimizer
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)
# Start training
logger.info("Training Started".center(60, '#'))
for epoch in range(epochs):
# Train the model
shuffle=True,
num_workers=16,
pin_memory=True)
validation_dataloader = DataLoader(val_data,
batch_size=batch_size,
shuffle=True,
num_workers=16,
pin_memory=True)
if args.model == 'r2plus1d_18':
model = cnn3d = r2plus1d_18(pretrained=False, num_classes=500)
elif args.model == 'r2plus1d_18_new':
model = r2plus1d_18_new(pretrained=False, progress=False)
elif args.model == '3dresnet18':
model = resnet18(pretrained=False,
sample_size=sample_size,
sample_duration=16,
num_classes=500)
elif args.model == 'ResCRNN':
model = ResCRNN(sample_size=112,
sample_duration=16,
num_classes=100,
arch="resnet18")
model = nn.DataParallel(model, device_ids=[0, 1, 2, 3], dim=0).to(device)
criterion = torch.nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=lr, weight_decay=0.0001)
scheduler = StepLR(optimizer, step_size=10, gamma=0.1)
best_epoch = 0
for epoch in tqdm(range(epoch_number)):
acc = train(model, device, training_dataloader, optimizer, epoch)
val_loss, val_acc = val(model, device, validation_dataloader)