utils.py

import torch
import torch.nn as nn

def type_tdouble(use_cuda=False):
    return torch.cuda.DoubleTensor if use_cuda else torch.DoubleTensor


def one_hot(labels, n_class, use_cuda=False):
    # Ensure labels are [N x 1]
    if len(list(labels.size())) == 1:
        labels = labels.unsqueeze(1)
    mask = type_tdouble(use_cuda)(labels.size(0), n_class).fill_(0)
    # scatter dimension, position indices, fill_value
    return mask.scatter_(1, labels, 1)


def init_weights(module):
    for m in module.modules():
        if isinstance(m, nn.Conv2d):
            nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
        elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
            nn.init.constant_(m.weight, 1)
            nn.init.constant_(m.bias, 0)





def execute_graph(model, loader, optimizer, scheduler, epoch, use_cuda):
    t_loss = train_validate(model, loader, optimizer, True, epoch, use_cuda)
    v_loss = train_validate(model, loader, optimizer, False, epoch, use_cuda)

    scheduler.step(v_loss)

    logger.add_scalar(log_dir + '/train-loss', t_loss, epoch)
    logger.add_scalar(log_dir + '/valid-loss', v_loss, epoch)

    return v_loss