diff --git a/losses.py b/losses.py index 1e9a080..ce05b0d 100755 --- a/losses.py +++ b/losses.py @@ -10,7 +10,19 @@ Shorthands for loss: - TripletLoss: htri - CenterLoss: cent """ -__all__ = ['CrossEntropyLabelSmooth', 'TripletLoss', 'CenterLoss'] +__all__ = ['DeepSupervision', 'CrossEntropyLabelSmooth', 'TripletLoss', 'CenterLoss'] + +def DeepSupervision(criterion, xs, y): + """ + Args: + criterion: loss function + xs: tuple of inputs + y: ground truth + """ + loss = 0. + for x in xs: + loss += criterion(x, y) + return loss class CrossEntropyLabelSmooth(nn.Module): """Cross entropy loss with label smoothing regularizer. diff --git a/models/HACNN.py b/models/HACNN.py index dc043eb..f5a832c 100644 --- a/models/HACNN.py +++ b/models/HACNN.py @@ -188,10 +188,11 @@ class HACNN(nn.Module): feat_dim (int): feature dimension for each branch learn_region (bool): whether to learn region features (i.e. local branch) """ - def __init__(self, num_classes, loss={'xent'}, nchannels=[128, 256, 384], feat_dim=512, learn_region=True, **kwargs): + def __init__(self, num_classes, loss={'xent'}, nchannels=[128, 256, 384], feat_dim=512, learn_region=True, use_gpu=True, **kwargs): super(HACNN, self).__init__() self.loss = loss self.learn_region = learn_region + self.use_gpu = use_gpu self.conv = ConvBlock(3, 32, 3, s=2, p=1) @@ -267,6 +268,7 @@ class HACNN(nn.Module): theta = torch.zeros(theta_i.size(0), 2, 3) theta[:,:,:2] = scale_factors theta[:,:,-1] = theta_i + if self.use_gpu: theta = theta.cuda() return theta def forward(self, x):