Skip to content

Instantly share code, notes, and snippets.

View paulomann's full-sized avatar

Paulo Mann paulomann

  • Rio de Janeiro, Brasil
View GitHub Profile
class Bottleneck(nn.Module):
''' Standard bottleneck block
input = inplanes * H * W
middle = planes * H/stride * W/stride
output = 4*planes * H/stride * W/stride
'''
expansion = 4
def __init__(self, inplanes, planes, stride=1, dilation=1, downsample=None):
super(Bottleneck, self).__init__()
class Trainer():
def __init__(self, model, dataloaders, dataset_sizes, criterion, optimizer,
scheduler, num_epochs=100, threshold=0.5):
self.acc_loss = {"train": {"loss": [], "acc": []},
"val": {"loss": [], "acc": []}}
self.device = torch.device(
"cuda:0" if torch.cuda.is_available() else "cpu")
self.model = model.to(self.device)
print("Using device ", self.device)