Kaggle competition: Severstal: Steel Defect Detection
<SOS> metric: Dice Kaggle data source: severstal:-steel-defect-detection<define_variables>
s_time = time.time()
warnings.filterwarnings('ignore')
class SAKTDataset(Dataset):
def __init__(self, group, n_skill, max_seq=100):
super(SAKTDataset, self ).__init__()
self.max_seq = max_seq
self.n_skill = n_skill
self.samples = group
self.user_ids = []
for user_id in group.index:
q, qa = group[user_id]
if len(q)< 10:
continue
self.user_ids.append(user_id)
 def __len__(self):
return len(self.user_ids)
 def __getitem__(self, index):
user_id = self.user_ids[index]
q_, qa_ = self.samples[user_id]
seq_len = len(q_)
 q = np.zeros(self.max_seq, dtype=int)
 qa = np.zeros(self.max_seq, dtype=int)
 if seq_len >= self.max_seq:
q[:] = q_[-self.max_seq:]
qa[:] = qa_[-self.max_seq:]
else:
q[-seq_len:] = q_
qa[-seq_len:] = qa_
target_id = q[1:]
label = qa[1:]
x = np.zeros(self.max_seq-1, dtype=int)
 x = q[:-1].copy()
x +=(qa[:-1] == 1)* self.n_skill
return x, target_id, label<split>
def post_process(mask, min_size):
num_component, component = cv2.connectedComponents(mask.astype(np.uint8))
predict = np.zeros(( 256, 1600), np.float32)
 for c in range(1, num_component):
p =(component == c)
 if p.sum() > min_size:
predict[p] = 1
return predict
def probability_mask_to_probability_label(probability):
probability = F.adaptive_max_pool2d(probability,1 ).squeeze(-1 ).squeeze(-1)
 return probability
def run_length_encode(mask):
m = mask.T.flatten()
if m.sum() ==0:
rle=''
else:
m = np.concatenate([[0], m, [0]])
 run = np.where(m[1:] != m[:-1])[0] + 1
run[1::2] -= run[::2]
rle = ' '.join(str(r)for r in run)
 return rle
train, val = train_test_split(group, test_size=0.2)
 train_dataset = SAKTDataset(train, n_skill)
 train_dataloader = DataLoader(train_dataset, batch_size=2048, shuffle=True, num_workers=8)
 del train
val_dataset = SAKTDataset(val, n_skill)
 val_dataloader = DataLoader(val_dataset, batch_size=2048, shuffle=True, num_workers=8)
 del val 
<define_search_model>
class SteelDataset(Dataset):
def __init__(self, defect_root,
resize_long=None, img_normal=True, hor_flip=False,
crop_size=None, to_torch=True, random_hue=False):
self.img_name_list = glob.glob(defect_root + '/*.jpg')
 self.defect_root = defect_root
self.resize_long = resize_long
self.crop_size = crop_size
self.img_normal = img_normal
self.hor_flip = hor_flip
self.to_torch = to_torch
self.random_hue = random_hue
def __len__(self):
return len(self.img_name_list)
 def __getitem__(self, idx):
name = self.img_name_list[idx]
img = Image.open(name)
 if self.resize_long:
target_long = random.randint(self.resize_long[0], self.resize_long[1])
 img = transforms.Resize(( 256, target_long))(img)
 if self.hor_flip:
img = transforms.RandomHorizontalFlip()(img)
 if self.crop_size:
img = transforms.RandomCrop(( self.crop_size[0], self.crop_size[1]), padding=0, pad_if_needed=True )(img)
 if self.random_hue:
img = transforms.ColorJitter(brightness=0.5, contrast=0.5, saturation=0.5, hue=0.25 )(img)
 if self.to_torch:
img = transforms.ToTensor()(img)
 if self.img_normal:
img = transforms.Normalize(mean=(0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225))(img)
 return {'name': name, 'img': img}
class SteelDatasetMSF(Dataset):
def __init__(self, defect_root,
resize_long=None, img_normal=True, hor_flip=False,
crop_size=None, to_torch=True, random_hue=False, scales=[1.]):
self.img_name_list = glob.glob(defect_root + '/*.jpg')
 self.defect_root = defect_root
self.scales = scales
self.resize_long = resize_long
self.crop_size = crop_size
self.img_normal = img_normal
self.hor_flip = hor_flip
self.to_torch = to_torch
self.random_hue = random_hue
def __len__(self):
return len(self.img_name_list)
 def __getitem__(self, idx):
name = self.img_name_list[idx]
img = Image.open(name)
 ori_w, ori_h = img.size
ms_img_list = []
for s in self.scales:
if s == 1:
s_img = img
else:
w = int(s*ori_w)
 s_img = transforms.Resize(( ori_h, w))(img)
 if self.random_hue:
s_img = transforms.ColorJitter(brightness=0.5, contrast=0.5, saturation=0.5, hue=0.25 )(s_img)
 if self.to_torch:
s_img = transforms.ToTensor()(s_img)
 if self.img_normal:
s_img = transforms.Normalize(mean=(0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225))(s_img)
 ms_img_list.append(s_img)
 if len(self.scales)== 1:
ms_img_list = ms_img_list[0]
return {'name': name, 'img': ms_img_list}
class FFN(nn.Module):
def __init__(self, state_size=200):
super(FFN, self ).__init__()
self.state_size = state_size
self.lr1 = nn.Linear(state_size, state_size)
 self.relu = nn.ReLU()
self.lr2 = nn.Linear(state_size, state_size)
 self.dropout = nn.Dropout(0.2)
 def forward(self, x):
x = self.lr1(x)
 x = self.relu(x)
 x = self.lr2(x)
 return self.dropout(x)
 def future_mask(seq_length):
future_mask = np.triu(np.ones(( seq_length, seq_length)) , k=1 ).astype('bool')
 return torch.from_numpy(future_mask)
 class SAKTModel(nn.Module):
def __init__(self, n_skill, max_seq=100, embed_dim=128):
super(SAKTModel, self ).__init__()
self.n_skill = n_skill
self.embed_dim = embed_dim
self.embedding = nn.Embedding(2*n_skill+1, embed_dim)
 self.pos_embedding = nn.Embedding(max_seq-1, embed_dim)
 self.e_embedding = nn.Embedding(n_skill+1, embed_dim)
 self.multi_att = nn.MultiheadAttention(embed_dim=embed_dim, num_heads=8, dropout=0.2)
 self.dropout = nn.Dropout(0.2)
 self.layer_normal = nn.LayerNorm(embed_dim)
 self.ffn = FFN(embed_dim)
 self.pred = nn.Linear(embed_dim, 1)
 def forward(self, x, question_ids):
device = x.device
x = self.embedding(x)
 pos_id = torch.arange(x.size(1)).unsqueeze(0 ).to(device)
 pos_x = self.pos_embedding(pos_id)
 x = x + pos_x
e = self.e_embedding(question_ids)
 x = x.permute(1, 0, 2)
 e = e.permute(1, 0, 2)
 att_mask = future_mask(x.size(0)).to(device)
 att_output, att_weight = self.multi_att(e, x, x, attn_mask=att_mask)
 att_output = self.layer_normal(att_output + e)
 att_output = att_output.permute(1, 0, 2)
 x = self.ffn(att_output)
 x = self.layer_normal(x + att_output)
 x = self.pred(x)
 return x.squeeze(-1), att_weight<choose_model_class>
BatchNorm2d = nn.BatchNorm2d
BN_MOMENTUM = 0.01
def conv3x3(in_planes, out_planes, stride=1):

 return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
padding=1, bias=False)
 class BasicBlock(nn.Module):
expansion = 1
def __init__(self, inplanes, planes, stride=1, downsample=None):
super(BasicBlock, self ).__init__()
self.conv1 = conv3x3(inplanes, planes, stride)
 self.bn1 = BatchNorm2d(planes, momentum=BN_MOMENTUM)
 self.relu = nn.ReLU(inplace=False)
 self.conv2 = conv3x3(planes, planes)
 self.bn2 = BatchNorm2d(planes, momentum=BN_MOMENTUM)
 self.downsample = downsample
self.stride = stride
def forward(self, x):
residual = x
out = self.conv1(x)
 out = self.bn1(out)
 out = self.relu(out)
 out = self.conv2(out)
 out = self.bn2(out)
 if self.downsample is not None:
residual = self.downsample(x)
 out = out + residual
out = self.relu(out)
 return out
class Bottleneck(nn.Module):
expansion = 4
def __init__(self, inplanes, planes, stride=1, downsample=None):
super(Bottleneck, self ).__init__()
self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
 self.bn1 = BatchNorm2d(planes, momentum=BN_MOMENTUM)
 self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
padding=1, bias=False)
 self.bn2 = BatchNorm2d(planes, momentum=BN_MOMENTUM)
 self.conv3 = nn.Conv2d(planes, planes * self.expansion, kernel_size=1,
bias=False)
 self.bn3 = BatchNorm2d(planes * self.expansion,
momentum=BN_MOMENTUM)
 self.relu = nn.ReLU(inplace=False)
 self.downsample = downsample
self.stride = stride
def forward(self, x):
residual = x
out = self.conv1(x)
 out = self.bn1(out)
 out = self.relu(out)
 out = self.conv2(out)
 out = self.bn2(out)
 out = self.relu(out)
 out = self.conv3(out)
 out = self.bn3(out)
 if self.downsample is not None:
residual = self.downsample(x)
 out = out + residual
out = self.relu(out)
 return out
class HighResolutionModule(nn.Module):
def __init__(self, num_branches, blocks, num_blocks, num_inchannels,
num_channels, fuse_method, multi_scale_output=True):
super(HighResolutionModule, self ).__init__()
self._check_branches(
 num_branches, blocks, num_blocks, num_inchannels, num_channels)
 self.num_inchannels = num_inchannels
self.fuse_method = fuse_method
self.num_branches = num_branches
self.multi_scale_output = multi_scale_output
self.branches = self._make_branches(
 num_branches, blocks, num_blocks, num_channels)
 self.fuse_layers = self._make_fuse_layers()
self.relu = nn.ReLU(inplace=False)
 def _check_branches(self, num_branches, blocks, num_blocks,
num_inchannels, num_channels):
if num_branches != len(num_blocks):
error_msg = 'NUM_BRANCHES({})<> NUM_BLOCKS({})'.format(
 num_branches, len(num_blocks))
logger.error(error_msg)
 raise ValueError(error_msg)
 if num_branches != len(num_channels):
error_msg = 'NUM_BRANCHES({})<> NUM_CHANNELS({})'.format(
 num_branches, len(num_channels))
logger.error(error_msg)
 raise ValueError(error_msg)
 if num_branches != len(num_inchannels):
error_msg = 'NUM_BRANCHES({})<> NUM_INCHANNELS({})'.format(
 num_branches, len(num_inchannels))
logger.error(error_msg)
 raise ValueError(error_msg)
 def _make_one_branch(self, branch_index, block, num_blocks, num_channels,
stride=1):
downsample = None
if stride != 1 or \
self.num_inchannels[branch_index] != num_channels[branch_index] * block.expansion:
downsample = nn.Sequential(
 nn.Conv2d(self.num_inchannels[branch_index],
num_channels[branch_index] * block.expansion,
kernel_size=1, stride=stride, bias=False),
BatchNorm2d(num_channels[branch_index] * block.expansion,
momentum=BN_MOMENTUM), 
)
 layers = []
layers.append(block(self.num_inchannels[branch_index],
num_channels[branch_index], stride, downsample))
self.num_inchannels[branch_index] = \
num_channels[branch_index] * block.expansion
for i in range(1, num_blocks[branch_index]):
layers.append(block(self.num_inchannels[branch_index],
num_channels[branch_index]))
return nn.Sequential(*layers)
 def _make_branches(self, num_branches, block, num_blocks, num_channels):
branches = []
for i in range(num_branches):
branches.append(
 self._make_one_branch(i, block, num_blocks, num_channels))
return nn.ModuleList(branches)
 def _make_fuse_layers(self):
if self.num_branches == 1:
return None
num_branches = self.num_branches
num_inchannels = self.num_inchannels
fuse_layers = []
for i in range(num_branches if self.multi_scale_output else 1):
fuse_layer = []
for j in range(num_branches):
if j > i:
fuse_layer.append(nn.Sequential(
 nn.Conv2d(num_inchannels[j],
num_inchannels[i],
1,
1,
0,
bias=False),
BatchNorm2d(num_inchannels[i], momentum=BN_MOMENTUM)))
 elif j == i:
fuse_layer.append(None)
 else:
conv3x3s = []
for k in range(i - j):
if k == i - j - 1:
num_outchannels_conv3x3 = num_inchannels[i]
conv3x3s.append(nn.Sequential(
 nn.Conv2d(num_inchannels[j],
num_outchannels_conv3x3,
3, 2, 1, bias=False),
BatchNorm2d(num_outchannels_conv3x3,
momentum=BN_MOMENTUM)))
 else:
num_outchannels_conv3x3 = num_inchannels[j]
conv3x3s.append(nn.Sequential(
 nn.Conv2d(num_inchannels[j],
num_outchannels_conv3x3,
3, 2, 1, bias=False),
BatchNorm2d(num_outchannels_conv3x3,
momentum=BN_MOMENTUM),
nn.ReLU(inplace=False)))
 fuse_layer.append(nn.Sequential(*conv3x3s))
fuse_layers.append(nn.ModuleList(fuse_layer))
return nn.ModuleList(fuse_layers)
 def get_num_inchannels(self):
return self.num_inchannels
def forward(self, x):
if self.num_branches == 1:
return [self.branches[0](x[0])]
for i in range(self.num_branches):
x[i] = self.branches[i](x[i])
 x_fuse = []
for i in range(len(self.fuse_layers)) :
y = x[0] if i == 0 else self.fuse_layers[i][0](x[0])
 for j in range(1, self.num_branches):
if i == j:
y = y + x[j]
elif j > i:
width_output = x[i].shape[-1]
height_output = x[i].shape[-2]
y = y + F.interpolate(
 self.fuse_layers[i][j](x[j]),
size=[height_output, width_output],
mode='bilinear')
 else:
y = y + self.fuse_layers[i][j](x[j])
 x_fuse.append(self.relu(y))
return x_fuse
blocks_dict = {
'BASIC': BasicBlock,
'BOTTLENECK': Bottleneck
}
class HighResolutionNet(nn.Module):
def __init__(self, config, **kwargs):
extra = config['MODEL']['EXTRA']
super(HighResolutionNet, self ).__init__()
self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=2, padding=1,
bias=False)
 self.bn1 = BatchNorm2d(64, momentum=BN_MOMENTUM)
 self.conv2 = nn.Conv2d(64, 64, kernel_size=3, stride=2, padding=1,
bias=False)
 self.bn2 = BatchNorm2d(64, momentum=BN_MOMENTUM)
 self.relu = nn.ReLU(inplace=False)
 self.stage1_cfg = extra['STAGE1']
num_channels = self.stage1_cfg['NUM_CHANNELS'][0]
block = blocks_dict[self.stage1_cfg['BLOCK']]
num_blocks = self.stage1_cfg['NUM_BLOCKS'][0]
self.layer1 = self._make_layer(block, 64, num_channels, num_blocks)
 stage1_out_channel = block.expansion * num_channels
self.stage2_cfg = extra['STAGE2']
num_channels = self.stage2_cfg['NUM_CHANNELS']
block = blocks_dict[self.stage2_cfg['BLOCK']]
num_channels = [
num_channels[i] * block.expansion for i in range(len(num_channels)) ]
self.transition1 = self._make_transition_layer(
 [stage1_out_channel], num_channels)
 self.stage2, pre_stage_channels = self._make_stage(
 self.stage2_cfg, num_channels)
 self.stage3_cfg = extra['STAGE3']
num_channels = self.stage3_cfg['NUM_CHANNELS']
block = blocks_dict[self.stage3_cfg['BLOCK']]
num_channels = [
num_channels[i] * block.expansion for i in range(len(num_channels)) ]
self.transition2 = self._make_transition_layer(
 pre_stage_channels, num_channels)
 self.stage3, pre_stage_channels = self._make_stage(
 self.stage3_cfg, num_channels)
 self.stage4_cfg = extra['STAGE4']
num_channels = self.stage4_cfg['NUM_CHANNELS']
block = blocks_dict[self.stage4_cfg['BLOCK']]
num_channels = [
num_channels[i] * block.expansion for i in range(len(num_channels)) ]
self.transition3 = self._make_transition_layer(
 pre_stage_channels, num_channels)
 self.stage4, pre_stage_channels = self._make_stage(
 self.stage4_cfg, num_channels, multi_scale_output=True)
 last_inp_channels = np.int(np.sum(pre_stage_channels))
self.last_layer = nn.Sequential(
 nn.Conv2d(
 in_channels=last_inp_channels,
out_channels=last_inp_channels,
kernel_size=1,
stride=1,
padding=0),
BatchNorm2d(last_inp_channels, momentum=BN_MOMENTUM),
nn.ReLU(inplace=False),
nn.Conv2d(
 in_channels=last_inp_channels,
out_channels=config['MODEL']['NUM_CLASSES'],
kernel_size=extra['FINAL_CONV_KERNEL'],
stride=1,
padding=1 if extra['FINAL_CONV_KERNEL'] == 3 else 0)
)
 self.backbone = nn.ModuleList([self.conv1, self.bn1, self.conv2, self.bn2, self.layer1, self.transition1,
self.stage2, self.transition2, self.stage3, self.transition3, self.stage4])
 self.fuse = nn.ModuleList([self.last_layer])
 def _make_transition_layer(
 self, num_channels_pre_layer, num_channels_cur_layer):
num_branches_cur = len(num_channels_cur_layer)
 num_branches_pre = len(num_channels_pre_layer)
 transition_layers = []
for i in range(num_branches_cur):
if i < num_branches_pre:
if num_channels_cur_layer[i] != num_channels_pre_layer[i]:
transition_layers.append(nn.Sequential(
 nn.Conv2d(num_channels_pre_layer[i],
num_channels_cur_layer[i],
3,
1,
1,
bias=False),
BatchNorm2d(
 num_channels_cur_layer[i], momentum=BN_MOMENTUM),
nn.ReLU(inplace=False)))
 else:
transition_layers.append(None)
 else:
conv3x3s = []
for j in range(i + 1 - num_branches_pre):
inchannels = num_channels_pre_layer[-1]
outchannels = num_channels_cur_layer[i] \
if j == i - num_branches_pre else inchannels
conv3x3s.append(nn.Sequential(
 nn.Conv2d(
 inchannels, outchannels, 3, 2, 1, bias=False),
BatchNorm2d(outchannels, momentum=BN_MOMENTUM),
nn.ReLU(inplace=False)))
 transition_layers.append(nn.Sequential(*conv3x3s))
return nn.ModuleList(transition_layers)
 def _make_layer(self, block, inplanes, planes, blocks, stride=1):
downsample = None
if stride != 1 or inplanes != planes * block.expansion:
downsample = nn.Sequential(
 nn.Conv2d(inplanes, planes * block.expansion,
kernel_size=1, stride=stride, bias=False),
BatchNorm2d(planes * block.expansion, momentum=BN_MOMENTUM), 
)
 layers = []
layers.append(block(inplanes, planes, stride, downsample))
inplanes = planes * block.expansion
for i in range(1, blocks):
layers.append(block(inplanes, planes))
return nn.Sequential(*layers)
 def _make_stage(self, layer_config, num_inchannels,
multi_scale_output=True):
num_modules = layer_config['NUM_MODULES']
num_branches = layer_config['NUM_BRANCHES']
num_blocks = layer_config['NUM_BLOCKS']
num_channels = layer_config['NUM_CHANNELS']
block = blocks_dict[layer_config['BLOCK']]
fuse_method = layer_config['FUSE_METHOD']
modules = []
for i in range(num_modules):
if not multi_scale_output and i == num_modules - 1:
reset_multi_scale_output = False
else:
reset_multi_scale_output = True
modules.append(
 HighResolutionModule(num_branches,
block,
num_blocks,
num_inchannels,
num_channels,
fuse_method,
reset_multi_scale_output)
)
 num_inchannels = modules[-1].get_num_inchannels()
return nn.Sequential(*modules), num_inchannels
def forward(self, x):
x = self.conv1(x)
 x = self.bn1(x)
 x = self.relu(x)
 x = self.conv2(x)
 x = self.bn2(x)
 x = self.relu(x)
 x = self.layer1(x)
 x_list = []
for i in range(self.stage2_cfg['NUM_BRANCHES']):
if self.transition1[i] is not None:
x_list.append(self.transition1[i](x))
else:
x_list.append(x)
 y_list = self.stage2(x_list)
 x_list = []
for i in range(self.stage3_cfg['NUM_BRANCHES']):
if self.transition2[i] is not None:
if i < self.stage2_cfg['NUM_BRANCHES']:
x_list.append(self.transition2[i](y_list[i]))
else:
x_list.append(self.transition2[i](y_list[-1]))
else:
x_list.append(y_list[i])
 y_list = self.stage3(x_list)
 x_list = []
for i in range(self.stage4_cfg['NUM_BRANCHES']):
if self.transition3[i] is not None:
if i < self.stage3_cfg['NUM_BRANCHES']:
x_list.append(self.transition3[i](y_list[i]))
else:
x_list.append(self.transition3[i](y_list[-1]))
else:
x_list.append(y_list[i])
 x = self.stage4(x_list)
 x0_h, x0_w = x[0].size(2), x[0].size(3)
 x1 = F.upsample(x[1], size=(x0_h, x0_w), mode='bilinear')
 x2 = F.upsample(x[2], size=(x0_h, x0_w), mode='bilinear')
 x3 = F.upsample(x[3], size=(x0_h, x0_w), mode='bilinear')
 x = torch.cat([x[0], x1, x2, x3], 1)
 x = self.last_layer(x)
 return x
def trainable_parameters(self):
return list(self.backbone.parameters()), list(self.fuse.parameters())
 def init_weights(self, pretrained='',):
for m in self.modules() :
if isinstance(m, nn.Conv2d):
nn.init.normal_(m.weight, std=0.001)
 elif isinstance(m, nn.BatchNorm2d):
nn.init.constant_(m.weight, 1)
 nn.init.constant_(m.bias, 0)
 if os.path.isfile(pretrained):
pretrained_dict = torch.load(pretrained)
 model_dict = self.state_dict()
pretrained_dict = {k: v for k, v in pretrained_dict.items()
if k in model_dict.keys() }
print('length of pretrained_dict:', len(pretrained_dict.items()))
model_dict.update(pretrained_dict)
 self.load_state_dict(model_dict)
 class CAM(HighResolutionNet):
def __init__(self, cfg, **kwargs):
super(CAM, self ).__init__(cfg, **kwargs)
 def forward(self, x):
x = self.conv1(x)
 x = self.bn1(x)
 x = self.relu(x)
 x = self.conv2(x)
 x = self.bn2(x)
 x = self.relu(x)
 x = self.layer1(x)
 x_list = []
for i in range(self.stage2_cfg['NUM_BRANCHES']):
if self.transition1[i] is not None:
x_list.append(self.transition1[i](x))
else:
x_list.append(x)
 y_list = self.stage2(x_list)
 x_list = []
for i in range(self.stage3_cfg['NUM_BRANCHES']):
if self.transition2[i] is not None:
if i < self.stage2_cfg['NUM_BRANCHES']:
x_list.append(self.transition2[i](y_list[i]))
else:
x_list.append(self.transition2[i](y_list[-1]))
else:
x_list.append(y_list[i])
 y_list = self.stage3(x_list)
 x_list = []
for i in range(self.stage4_cfg['NUM_BRANCHES']):
if self.transition3[i] is not None:
if i < self.stage3_cfg['NUM_BRANCHES']:
x_list.append(self.transition3[i](y_list[i]))
else:
x_list.append(self.transition3[i](y_list[-1]))
else:
x_list.append(y_list[i])
 x = self.stage4(x_list)
 x0_h, x0_w = x[0].size(2), x[0].size(3)
 x1 = F.upsample(x[1], size=(x0_h, x0_w), mode='bilinear')
 x2 = F.upsample(x[2], size=(x0_h, x0_w), mode='bilinear')
 x3 = F.upsample(x[3], size=(x0_h, x0_w), mode='bilinear')
 x = torch.cat([x[0], x1, x2, x3], 1)
 x = self.last_layer(x)
 x = F.relu(x)
 x = x[0] + x[1].flip(-1)
 return x
def get_hr_seg_model(cfg, **kwargs):
cfg['MODEL']['NUM_CLASSES'] = 5
model = HighResolutionNet(cfg, **kwargs)
 return model
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model = SAKTModel(n_skill, embed_dim=128)
 optimizer = torch.optim.Adam(model.parameters() , lr=1e-3)
 criterion = nn.BCEWithLogitsLoss()
model.to(device)
 criterion.to(device )<train_model>
ALIGN_CORNERS = True
BN_MOMENTUM = 0.1
relu_inplace = True
BatchNorm2d_class = BatchNorm2d = torch.nn.BatchNorm2d
class ModuleHelper:
@staticmethod
def BNReLU(num_features, bn_type=None, **kwargs):
return nn.Sequential(
 BatchNorm2d(num_features, **kwargs),
nn.ReLU() 
)
 @staticmethod
def BatchNorm2d(*args, **kwargs):
return BatchNorm2d
def conv3x3(in_planes, out_planes, stride=1):

 return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
padding=1, bias=False)
 class SpatialGather_Module(nn.Module):

 def __init__(self, cls_num=0, scale=1):
super(SpatialGather_Module, self ).__init__()
self.cls_num = cls_num
self.scale = scale
def forward(self, feats, probs):
batch_size, c, h, w = probs.size(0), probs.size(1), probs.size(2), probs.size(3)
 probs = probs.view(batch_size, c, -1)
 feats = feats.view(batch_size, feats.size(1), -1)
 feats = feats.permute(0, 2, 1)
 probs = F.softmax(self.scale * probs, dim=2)
 ocr_context = torch.matmul(probs, feats)\ 
.permute(0, 2, 1 ).unsqueeze(3)
 return ocr_context
class _ObjectAttentionBlock(nn.Module):

 def __init__(self,
in_channels,
key_channels,
scale=1,
bn_type=None):
super(_ObjectAttentionBlock, self ).__init__()
self.scale = scale
self.in_channels = in_channels
self.key_channels = key_channels
self.pool = nn.MaxPool2d(kernel_size=(scale, scale))
self.f_pixel = nn.Sequential(
 nn.Conv2d(in_channels=self.in_channels, out_channels=self.key_channels,
kernel_size=1, stride=1, padding=0, bias=False),
ModuleHelper.BNReLU(self.key_channels, bn_type=bn_type),
nn.Conv2d(in_channels=self.key_channels, out_channels=self.key_channels,
kernel_size=1, stride=1, padding=0, bias=False),
ModuleHelper.BNReLU(self.key_channels, bn_type=bn_type), 
)
 self.f_object = nn.Sequential(
 nn.Conv2d(in_channels=self.in_channels, out_channels=self.key_channels,
kernel_size=1, stride=1, padding=0, bias=False),
ModuleHelper.BNReLU(self.key_channels, bn_type=bn_type),
nn.Conv2d(in_channels=self.key_channels, out_channels=self.key_channels,
kernel_size=1, stride=1, padding=0, bias=False),
ModuleHelper.BNReLU(self.key_channels, bn_type=bn_type), 
)
 self.f_down = nn.Sequential(
 nn.Conv2d(in_channels=self.in_channels, out_channels=self.key_channels,
kernel_size=1, stride=1, padding=0, bias=False),
ModuleHelper.BNReLU(self.key_channels, bn_type=bn_type), 
)
 self.f_up = nn.Sequential(
 nn.Conv2d(in_channels=self.key_channels, out_channels=self.in_channels,
kernel_size=1, stride=1, padding=0, bias=False),
ModuleHelper.BNReLU(self.in_channels, bn_type=bn_type), 
)
 def forward(self, x, proxy):
batch_size, h, w = x.size(0), x.size(2), x.size(3)
 if self.scale > 1:
x = self.pool(x)
 query = self.f_pixel(x ).view(batch_size, self.key_channels, -1)
 query = query.permute(0, 2, 1)
 key = self.f_object(proxy ).view(batch_size, self.key_channels, -1)
 value = self.f_down(proxy ).view(batch_size, self.key_channels, -1)
 value = value.permute(0, 2, 1)
 sim_map = torch.matmul(query, key)
 sim_map =(self.key_channels**-.5)* sim_map
sim_map = F.softmax(sim_map, dim=-1)
 context = torch.matmul(sim_map, value)
 context = context.permute(0, 2, 1 ).contiguous()
context = context.view(batch_size, self.key_channels, *x.size() [2:])
 context = self.f_up(context)
 if self.scale > 1:
context = F.interpolate(input=context, size=(h, w), mode='bilinear', align_corners=ALIGN_CORNERS)
 return context
class ObjectAttentionBlock2D(_ObjectAttentionBlock):
def __init__(self,
in_channels,
key_channels,
scale=1,
bn_type=None):
super(ObjectAttentionBlock2D, self ).__init__(in_channels,
key_channels,
scale,
bn_type=bn_type)
 class SpatialOCR_Module(nn.Module):

 def __init__(self,
in_channels,
key_channels,
out_channels,
scale=1,
dropout=0.1,
bn_type=None):
super(SpatialOCR_Module, self ).__init__()
self.object_context_block = ObjectAttentionBlock2D(in_channels,
key_channels,
scale,
bn_type)
 _in_channels = 2 * in_channels
self.conv_bn_dropout = nn.Sequential(
 nn.Conv2d(_in_channels, out_channels, kernel_size=1, padding=0, bias=False),
ModuleHelper.BNReLU(out_channels, bn_type=bn_type),
nn.Dropout2d(dropout)
)
 def forward(self, feats, proxy_feats):
context = self.object_context_block(feats, proxy_feats)
 output = self.conv_bn_dropout(torch.cat([context, feats], 1))
return output
class BasicBlock(nn.Module):
expansion = 1
def __init__(self, inplanes, planes, stride=1, downsample=None):
super(BasicBlock, self ).__init__()
self.conv1 = conv3x3(inplanes, planes, stride)
 self.bn1 = BatchNorm2d(planes, momentum=BN_MOMENTUM)
 self.relu = nn.ReLU(inplace=True)
 self.conv2 = conv3x3(planes, planes)
 self.bn2 = BatchNorm2d(planes, momentum=BN_MOMENTUM)
 self.downsample = downsample
self.stride = stride
def forward(self, x):
residual = x
out = self.conv1(x)
 out = self.bn1(out)
 out = self.relu(out)
 out = self.conv2(out)
 out = self.bn2(out)
 if self.downsample is not None:
residual = self.downsample(x)
 out = out + residual
out = self.relu(out)
 return out
class Bottleneck(nn.Module):
expansion = 4
def __init__(self, inplanes, planes, stride=1, downsample=None):
super(Bottleneck, self ).__init__()
self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
 self.bn1 = BatchNorm2d(planes, momentum=BN_MOMENTUM)
 self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
padding=1, bias=False)
 self.bn2 = BatchNorm2d(planes, momentum=BN_MOMENTUM)
 self.conv3 = nn.Conv2d(planes, planes * self.expansion, kernel_size=1,
bias=False)
 self.bn3 = BatchNorm2d(planes * self.expansion,
momentum=BN_MOMENTUM)
 self.relu = nn.ReLU(inplace=relu_inplace)
 self.downsample = downsample
self.stride = stride
def forward(self, x):
residual = x
out = self.conv1(x)
 out = self.bn1(out)
 out = self.relu(out)
 out = self.conv2(out)
 out = self.bn2(out)
 out = self.relu(out)
 out = self.conv3(out)
 out = self.bn3(out)
 if self.downsample is not None:
residual = self.downsample(x)
 out = out + residual
out = self.relu(out)
 return out
class HighResolutionModule(nn.Module):
def __init__(self, num_branches, blocks, num_blocks, num_inchannels,
num_channels, fuse_method, multi_scale_output=True):
super(HighResolutionModule, self ).__init__()
self._check_branches(
 num_branches, blocks, num_blocks, num_inchannels, num_channels)
 self.num_inchannels = num_inchannels
self.fuse_method = fuse_method
self.num_branches = num_branches
self.multi_scale_output = multi_scale_output
self.branches = self._make_branches(
 num_branches, blocks, num_blocks, num_channels)
 self.fuse_layers = self._make_fuse_layers()
self.relu = nn.ReLU(inplace=relu_inplace)
 def _check_branches(self, num_branches, blocks, num_blocks,
num_inchannels, num_channels):
if num_branches != len(num_blocks):
error_msg = 'NUM_BRANCHES({})<> NUM_BLOCKS({})'.format(
 num_branches, len(num_blocks))
raise ValueError(error_msg)
 if num_branches != len(num_channels):
error_msg = 'NUM_BRANCHES({})<> NUM_CHANNELS({})'.format(
 num_branches, len(num_channels))
raise ValueError(error_msg)
 if num_branches != len(num_inchannels):
error_msg = 'NUM_BRANCHES({})<> NUM_INCHANNELS({})'.format(
 num_branches, len(num_inchannels))
raise ValueError(error_msg)
 def _make_one_branch(self, branch_index, block, num_blocks, num_channels,
stride=1):
downsample = None
if stride != 1 or \
self.num_inchannels[branch_index] != num_channels[branch_index] * block.expansion:
downsample = nn.Sequential(
 nn.Conv2d(self.num_inchannels[branch_index],
num_channels[branch_index] * block.expansion,
kernel_size=1, stride=stride, bias=False),
BatchNorm2d(num_channels[branch_index] * block.expansion,
momentum=BN_MOMENTUM), 
)
 layers = []
layers.append(block(self.num_inchannels[branch_index],
num_channels[branch_index], stride, downsample))
self.num_inchannels[branch_index] = \
num_channels[branch_index] * block.expansion
for i in range(1, num_blocks[branch_index]):
layers.append(block(self.num_inchannels[branch_index],
num_channels[branch_index]))
return nn.Sequential(*layers)
 def _make_branches(self, num_branches, block, num_blocks, num_channels):
branches = []
for i in range(num_branches):
branches.append(
 self._make_one_branch(i, block, num_blocks, num_channels))
return nn.ModuleList(branches)
 def _make_fuse_layers(self):
if self.num_branches == 1:
return None
num_branches = self.num_branches
num_inchannels = self.num_inchannels
fuse_layers = []
for i in range(num_branches if self.multi_scale_output else 1):
fuse_layer = []
for j in range(num_branches):
if j > i:
fuse_layer.append(nn.Sequential(
 nn.Conv2d(num_inchannels[j],
num_inchannels[i],
1,
1,
0,
bias=False),
BatchNorm2d(num_inchannels[i], momentum=BN_MOMENTUM)))
 elif j == i:
fuse_layer.append(None)
 else:
conv3x3s = []
for k in range(i-j):
if k == i - j - 1:
num_outchannels_conv3x3 = num_inchannels[i]
conv3x3s.append(nn.Sequential(
 nn.Conv2d(num_inchannels[j],
num_outchannels_conv3x3,
3, 2, 1, bias=False),
BatchNorm2d(num_outchannels_conv3x3,
momentum=BN_MOMENTUM)))
 else:
num_outchannels_conv3x3 = num_inchannels[j]
conv3x3s.append(nn.Sequential(
 nn.Conv2d(num_inchannels[j],
num_outchannels_conv3x3,
3, 2, 1, bias=False),
BatchNorm2d(num_outchannels_conv3x3,
momentum=BN_MOMENTUM),
nn.ReLU(inplace=relu_inplace)))
 fuse_layer.append(nn.Sequential(*conv3x3s))
fuse_layers.append(nn.ModuleList(fuse_layer))
return nn.ModuleList(fuse_layers)
 def get_num_inchannels(self):
return self.num_inchannels
def forward(self, x):
if self.num_branches == 1:
return [self.branches[0](x[0])]
for i in range(self.num_branches):
x[i] = self.branches[i](x[i])
 x_fuse = []
for i in range(len(self.fuse_layers)) :
y = x[0] if i == 0 else self.fuse_layers[i][0](x[0])
 for j in range(1, self.num_branches):
if i == j:
y = y + x[j]
elif j > i:
width_output = x[i].shape[-1]
height_output = x[i].shape[-2]
y = y + F.interpolate(
 self.fuse_layers[i][j](x[j]),
size=[height_output, width_output],
mode='bilinear', align_corners=ALIGN_CORNERS)
 else:
y = y + self.fuse_layers[i][j](x[j])
 x_fuse.append(self.relu(y))
return x_fuse
blocks_dict = {
'BASIC': BasicBlock,
'BOTTLENECK': Bottleneck
}
class HighResolutionOCRNet(nn.Module):
def __init__(self, config, **kwargs):
global ALIGN_CORNERS
extra = config['MODEL']['EXTRA']
ocr_cfg = config['MODEL']['OCR']
super(HighResolutionOCRNet, self ).__init__()
ALIGN_CORNERS = True
self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=2, padding=1,
bias=False)
 self.bn1 = BatchNorm2d(64, momentum=BN_MOMENTUM)
 self.conv2 = nn.Conv2d(64, 64, kernel_size=3, stride=2, padding=1,
bias=False)
 self.bn2 = BatchNorm2d(64, momentum=BN_MOMENTUM)
 self.relu = nn.ReLU(inplace=relu_inplace)
 self.stage1_cfg = extra['STAGE1']
num_channels = self.stage1_cfg['NUM_CHANNELS'][0]
block = blocks_dict[self.stage1_cfg['BLOCK']]
num_blocks = self.stage1_cfg['NUM_BLOCKS'][0]
self.layer1 = self._make_layer(block, 64, num_channels, num_blocks)
 stage1_out_channel = block.expansion*num_channels
self.stage2_cfg = extra['STAGE2']
num_channels = self.stage2_cfg['NUM_CHANNELS']
block = blocks_dict[self.stage2_cfg['BLOCK']]
num_channels = [
num_channels[i] * block.expansion for i in range(len(num_channels)) ]
self.transition1 = self._make_transition_layer(
 [stage1_out_channel], num_channels)
 self.stage2, pre_stage_channels = self._make_stage(
 self.stage2_cfg, num_channels)
 self.stage3_cfg = extra['STAGE3']
num_channels = self.stage3_cfg['NUM_CHANNELS']
block = blocks_dict[self.stage3_cfg['BLOCK']]
num_channels = [
num_channels[i] * block.expansion for i in range(len(num_channels)) ]
self.transition2 = self._make_transition_layer(
 pre_stage_channels, num_channels)
 self.stage3, pre_stage_channels = self._make_stage(
 self.stage3_cfg, num_channels)
 self.stage4_cfg = extra['STAGE4']
num_channels = self.stage4_cfg['NUM_CHANNELS']
block = blocks_dict[self.stage4_cfg['BLOCK']]
num_channels = [
num_channels[i] * block.expansion for i in range(len(num_channels)) ]
self.transition3 = self._make_transition_layer(
 pre_stage_channels, num_channels)
 self.stage4, pre_stage_channels = self._make_stage(
 self.stage4_cfg, num_channels, multi_scale_output=True)
 last_inp_channels = np.int(np.sum(pre_stage_channels))
ocr_mid_channels = ocr_cfg['MID_CHANNELS']
ocr_key_channels = ocr_cfg['KEY_CHANNELS']
self.conv3x3_ocr = nn.Sequential(
 nn.Conv2d(last_inp_channels, ocr_mid_channels,
kernel_size=3, stride=1, padding=1),
BatchNorm2d(ocr_mid_channels),
nn.ReLU(inplace=relu_inplace), 
)
 self.ocr_gather_head = SpatialGather_Module(config['DATASET']['NUM_CLASSES'])
 self.ocr_distri_head = SpatialOCR_Module(in_channels=ocr_mid_channels,
key_channels=ocr_key_channels,
out_channels=ocr_mid_channels,
scale=1,
dropout=0.05, 
)
 self.cls_head = nn.Conv2d(
 ocr_mid_channels, config['DATASET']['NUM_CLASSES'], kernel_size=1, stride=1, padding=0, bias=True)
 self.aux_head = nn.Sequential(
 nn.Conv2d(last_inp_channels, last_inp_channels,
kernel_size=1, stride=1, padding=0),
BatchNorm2d(last_inp_channels),
nn.ReLU(inplace=relu_inplace),
nn.Conv2d(last_inp_channels, config['DATASET']['NUM_CLASSES'],
kernel_size=1, stride=1, padding=0, bias=True)
)
 self.backbone = nn.ModuleList([self.conv1, self.bn1, self.conv2, self.bn2, self.layer1, self.transition1,
self.stage2, self.transition2, self.stage3, self.transition3, self.stage4])
 self.fuse = nn.ModuleList([self.aux_head, self.conv3x3_ocr, self.ocr_gather_head, self.ocr_distri_head, self.cls_head])
 def _make_transition_layer(
 self, num_channels_pre_layer, num_channels_cur_layer):
num_branches_cur = len(num_channels_cur_layer)
 num_branches_pre = len(num_channels_pre_layer)
 transition_layers = []
for i in range(num_branches_cur):
if i < num_branches_pre:
if num_channels_cur_layer[i] != num_channels_pre_layer[i]:
transition_layers.append(nn.Sequential(
 nn.Conv2d(num_channels_pre_layer[i],
num_channels_cur_layer[i],
3,
1,
1,
bias=False),
BatchNorm2d(
 num_channels_cur_layer[i], momentum=BN_MOMENTUM),
nn.ReLU(inplace=relu_inplace)))
 else:
transition_layers.append(None)
 else:
conv3x3s = []
for j in range(i+1-num_branches_pre):
inchannels = num_channels_pre_layer[-1]
outchannels = num_channels_cur_layer[i] \
if j == i-num_branches_pre else inchannels
conv3x3s.append(nn.Sequential(
 nn.Conv2d(
 inchannels, outchannels, 3, 2, 1, bias=False),
BatchNorm2d(outchannels, momentum=BN_MOMENTUM),
nn.ReLU(inplace=relu_inplace)))
 transition_layers.append(nn.Sequential(*conv3x3s))
return nn.ModuleList(transition_layers)
 def _make_layer(self, block, inplanes, planes, blocks, stride=1):
downsample = None
if stride != 1 or inplanes != planes * block.expansion:
downsample = nn.Sequential(
 nn.Conv2d(inplanes, planes * block.expansion,
kernel_size=1, stride=stride, bias=False),
BatchNorm2d(planes * block.expansion, momentum=BN_MOMENTUM), 
)
 layers = []
layers.append(block(inplanes, planes, stride, downsample))
inplanes = planes * block.expansion
for i in range(1, blocks):
layers.append(block(inplanes, planes))
return nn.Sequential(*layers)
 def _make_stage(self, layer_config, num_inchannels,
multi_scale_output=True):
num_modules = layer_config['NUM_MODULES']
num_branches = layer_config['NUM_BRANCHES']
num_blocks = layer_config['NUM_BLOCKS']
num_channels = layer_config['NUM_CHANNELS']
block = blocks_dict[layer_config['BLOCK']]
fuse_method = layer_config['FUSE_METHOD']
modules = []
for i in range(num_modules):
if not multi_scale_output and i == num_modules - 1:
reset_multi_scale_output = False
else:
reset_multi_scale_output = True
modules.append(
 HighResolutionModule(num_branches,
block,
num_blocks,
num_inchannels,
num_channels,
fuse_method,
reset_multi_scale_output)
)
 num_inchannels = modules[-1].get_num_inchannels()
return nn.Sequential(*modules), num_inchannels
def forward(self, x):
x = self.conv1(x)
 x = self.bn1(x)
 x = self.relu(x)
 x = self.conv2(x)
 x = self.bn2(x)
 x = self.relu(x)
 x = self.layer1(x)
 x_list = []
for i in range(self.stage2_cfg['NUM_BRANCHES']):
if self.transition1[i] is not None:
x_list.append(self.transition1[i](x))
else:
x_list.append(x)
 y_list = self.stage2(x_list)
 x_list = []
for i in range(self.stage3_cfg['NUM_BRANCHES']):
if self.transition2[i] is not None:
if i < self.stage2_cfg['NUM_BRANCHES']:
x_list.append(self.transition2[i](y_list[i]))
else:
x_list.append(self.transition2[i](y_list[-1]))
else:
x_list.append(y_list[i])
 y_list = self.stage3(x_list)
 x_list = []
for i in range(self.stage4_cfg['NUM_BRANCHES']):
if self.transition3[i] is not None:
if i < self.stage3_cfg['NUM_BRANCHES']:
x_list.append(self.transition3[i](y_list[i]))
else:
x_list.append(self.transition3[i](y_list[-1]))
else:
x_list.append(y_list[i])
 x = self.stage4(x_list)
 x0_h, x0_w = x[0].size(2), x[0].size(3)
 x1 = F.interpolate(x[1], size=(x0_h, x0_w),
mode='bilinear', align_corners=ALIGN_CORNERS)
 x2 = F.interpolate(x[2], size=(x0_h, x0_w),
mode='bilinear', align_corners=ALIGN_CORNERS)
 x3 = F.interpolate(x[3], size=(x0_h, x0_w),
mode='bilinear', align_corners=ALIGN_CORNERS)
 feats = torch.cat([x[0], x1, x2, x3], 1)
 out_aux_seg = []
out_aux = self.aux_head(feats)
 feats = self.conv3x3_ocr(feats)
 context = self.ocr_gather_head(feats, out_aux)
 feats = self.ocr_distri_head(feats, context)
 out = self.cls_head(feats)
 out_aux_seg.append(out_aux)
 out_aux_seg.append(out)
 return out_aux_seg
def init_weights(self, pretrained='',):
print('=> init weights from normal distribution')
 for name, m in self.named_modules() :
if any(part in name for part in {'cls', 'aux', 'ocr'}):
continue
if isinstance(m, nn.Conv2d):
nn.init.normal_(m.weight, std=0.001)
 elif isinstance(m, BatchNorm2d):
nn.init.constant_(m.weight, 1)
 nn.init.constant_(m.bias, 0)
 if os.path.isfile(pretrained):
pretrained_dict = torch.load(pretrained, map_location={'cuda:0': 'cpu'})
 print('=> loading pretrained model {}'.format(pretrained))
model_dict = self.state_dict()
pretrained_dict = {k.replace('last_layer', 'aux_head' ).replace('model.', ''): v for k, v in pretrained_dict.items() }
pretrained_dict = {k: v for k, v in pretrained_dict.items()
if k in model_dict.keys() }
print('length of pretrained_dict:', len(pretrained_dict.items()))
model_dict.update(pretrained_dict)
 self.load_state_dict(model_dict)
 elif pretrained:
raise RuntimeError('No such file {}'.format(pretrained))
def trainable_parameters(self):
return list(self.backbone.parameters()), list(self.fuse.parameters())
 def get_hr_ocr_seg_model(cfg, **kwargs):
cfg['MODEL']['NUM_CLASSES'] = kwargs['num_classes']
model = HighResolutionOCRNet(cfg, **kwargs)
 return model
def train_epoch(model, train_iterator, optim, criterion, device="cpu"):
model.train()
train_loss = []
num_corrects = 0
num_total = 0
labels = []
outs = []
tbar = tqdm(train_iterator)
 for item in tbar:
x = item[0].to(device ).long()
target_id = item[1].to(device ).long()
label = item[2].to(device ).float()
target_mask =(target_id != 0)
 optim.zero_grad()
output, atten_weight = model(x, target_id)
 output = torch.masked_select(output, target_mask)
 label = torch.masked_select(label, target_mask)
 loss = criterion(output, label)
 loss.backward()
optim.step()
train_loss.append(loss.item())
 pred =(torch.sigmoid(output)>= 0.5 ).long()
num_corrects +=(pred == label ).sum().item()
num_total += len(label)
 labels.extend(label.view(-1 ).data.cpu().numpy())
 outs.extend(output.view(-1 ).data.cpu().numpy())
 tbar.set_description('loss - {:.4f}'.format(loss))
acc = num_corrects / num_total
auc = roc_auc_score(labels, outs)
 loss = np.average(train_loss)
 return loss, acc, auc<init_hyperparams>
ALIGN_CORNERS = True
BN_MOMENTUM = 0.1
relu_inplace = True
BatchNorm2d_class = BatchNorm2d = torch.nn.BatchNorm2d
class ModuleHelper:
@staticmethod
def BNReLU(num_features, bn_type=None, **kwargs):
return nn.Sequential(
 BatchNorm2d(num_features, **kwargs),
nn.ReLU() 
)
 @staticmethod
def BatchNorm2d(*args, **kwargs):
return BatchNorm2d
def conv3x3(in_planes, out_planes, stride=1):

 return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
padding=1, bias=False)
 class SpatialGather_Module(nn.Module):

 def __init__(self, cls_num=0, scale=1):
super(SpatialGather_Module, self ).__init__()
self.cls_num = cls_num
self.scale = scale
def forward(self, feats, probs):
batch_size, c, h, w = probs.size(0), probs.size(1), probs.size(2), probs.size(3)
 probs = probs.view(batch_size, c, -1)
 feats = feats.view(batch_size, feats.size(1), -1)
 feats = feats.permute(0, 2, 1)
 probs = F.softmax(self.scale * probs, dim=2)
 ocr_context = torch.matmul(probs, feats)\ 
.permute(0, 2, 1 ).unsqueeze(3)
 return ocr_context
class _ObjectAttentionBlock(nn.Module):

 def __init__(self,
in_channels,
key_channels,
scale=1,
bn_type=None):
super(_ObjectAttentionBlock, self ).__init__()
self.scale = scale
self.in_channels = in_channels
self.key_channels = key_channels
self.pool = nn.MaxPool2d(kernel_size=(scale, scale))
self.f_pixel = nn.Sequential(
 nn.Conv2d(in_channels=self.in_channels, out_channels=self.key_channels,
kernel_size=1, stride=1, padding=0, bias=False),
ModuleHelper.BNReLU(self.key_channels, bn_type=bn_type),
nn.Conv2d(in_channels=self.key_channels, out_channels=self.key_channels,
kernel_size=1, stride=1, padding=0, bias=False),
ModuleHelper.BNReLU(self.key_channels, bn_type=bn_type), 
)
 self.f_object = nn.Sequential(
 nn.Conv2d(in_channels=self.in_channels, out_channels=self.key_channels,
kernel_size=1, stride=1, padding=0, bias=False),
ModuleHelper.BNReLU(self.key_channels, bn_type=bn_type),
nn.Conv2d(in_channels=self.key_channels, out_channels=self.key_channels,
kernel_size=1, stride=1, padding=0, bias=False),
ModuleHelper.BNReLU(self.key_channels, bn_type=bn_type), 
)
 self.f_down = nn.Sequential(
 nn.Conv2d(in_channels=self.in_channels, out_channels=self.key_channels,
kernel_size=1, stride=1, padding=0, bias=False),
ModuleHelper.BNReLU(self.key_channels, bn_type=bn_type), 
)
 self.f_up = nn.Sequential(
 nn.Conv2d(in_channels=self.key_channels, out_channels=self.in_channels,
kernel_size=1, stride=1, padding=0, bias=False),
ModuleHelper.BNReLU(self.in_channels, bn_type=bn_type), 
)
 def forward(self, x, proxy):
batch_size, h, w = x.size(0), x.size(2), x.size(3)
 if self.scale > 1:
x = self.pool(x)
 query = self.f_pixel(x ).view(batch_size, self.key_channels, -1)
 query = query.permute(0, 2, 1)
 key = self.f_object(proxy ).view(batch_size, self.key_channels, -1)
 value = self.f_down(proxy ).view(batch_size, self.key_channels, -1)
 value = value.permute(0, 2, 1)
 sim_map = torch.matmul(query, key)
 sim_map =(self.key_channels**-.5)* sim_map
sim_map = F.softmax(sim_map, dim=-1)
 context = torch.matmul(sim_map, value)
 context = context.permute(0, 2, 1 ).contiguous()
context = context.view(batch_size, self.key_channels, *x.size() [2:])
 context = self.f_up(context)
 if self.scale > 1:
context = F.interpolate(input=context, size=(h, w), mode='bilinear', align_corners=ALIGN_CORNERS)
 return context
class ObjectAttentionBlock2D(_ObjectAttentionBlock):
def __init__(self,
in_channels,
key_channels,
scale=1,
bn_type=None):
super(ObjectAttentionBlock2D, self ).__init__(in_channels,
key_channels,
scale,
bn_type=bn_type)
 class SpatialOCR_Module(nn.Module):

 def __init__(self,
in_channels,
key_channels,
out_channels,
scale=1,
dropout=0.1,
bn_type=None):
super(SpatialOCR_Module, self ).__init__()
self.object_context_block = ObjectAttentionBlock2D(in_channels,
key_channels,
scale,
bn_type)
 _in_channels = 2 * in_channels
self.conv_bn_dropout = nn.Sequential(
 nn.Conv2d(_in_channels, out_channels, kernel_size=1, padding=0, bias=False),
ModuleHelper.BNReLU(out_channels, bn_type=bn_type),
nn.Dropout2d(dropout)
)
 def forward(self, feats, proxy_feats):
context = self.object_context_block(feats, proxy_feats)
 output = self.conv_bn_dropout(torch.cat([context, feats], 1))
return output
class BasicBlock(nn.Module):
expansion = 1
def __init__(self, inplanes, planes, stride=1, downsample=None):
super(BasicBlock, self ).__init__()
self.conv1 = conv3x3(inplanes, planes, stride)
 self.bn1 = BatchNorm2d(planes, momentum=BN_MOMENTUM)
 self.relu = nn.ReLU(inplace=True)
 self.conv2 = conv3x3(planes, planes)
 self.bn2 = BatchNorm2d(planes, momentum=BN_MOMENTUM)
 self.downsample = downsample
self.stride = stride
def forward(self, x):
residual = x
out = self.conv1(x)
 out = self.bn1(out)
 out = self.relu(out)
 out = self.conv2(out)
 out = self.bn2(out)
 if self.downsample is not None:
residual = self.downsample(x)
 out = out + residual
out = self.relu(out)
 return out
class Bottleneck(nn.Module):
expansion = 4
def __init__(self, inplanes, planes, stride=1, downsample=None):
super(Bottleneck, self ).__init__()
self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
 self.bn1 = BatchNorm2d(planes, momentum=BN_MOMENTUM)
 self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
padding=1, bias=False)
 self.bn2 = BatchNorm2d(planes, momentum=BN_MOMENTUM)
 self.conv3 = nn.Conv2d(planes, planes * self.expansion, kernel_size=1,
bias=False)
 self.bn3 = BatchNorm2d(planes * self.expansion,
momentum=BN_MOMENTUM)
 self.relu = nn.ReLU(inplace=relu_inplace)
 self.downsample = downsample
self.stride = stride
def forward(self, x):
residual = x
out = self.conv1(x)
 out = self.bn1(out)
 out = self.relu(out)
 out = self.conv2(out)
 out = self.bn2(out)
 out = self.relu(out)
 out = self.conv3(out)
 out = self.bn3(out)
 if self.downsample is not None:
residual = self.downsample(x)
 out = out + residual
out = self.relu(out)
 return out
class HighResolutionModule(nn.Module):
def __init__(self, num_branches, blocks, num_blocks, num_inchannels,
num_channels, fuse_method, multi_scale_output=True):
super(HighResolutionModule, self ).__init__()
self._check_branches(
 num_branches, blocks, num_blocks, num_inchannels, num_channels)
 self.num_inchannels = num_inchannels
self.fuse_method = fuse_method
self.num_branches = num_branches
self.multi_scale_output = multi_scale_output
self.branches = self._make_branches(
 num_branches, blocks, num_blocks, num_channels)
 self.fuse_layers = self._make_fuse_layers()
self.relu = nn.ReLU(inplace=relu_inplace)
 def _check_branches(self, num_branches, blocks, num_blocks,
num_inchannels, num_channels):
if num_branches != len(num_blocks):
error_msg = 'NUM_BRANCHES({})<> NUM_BLOCKS({})'.format(
 num_branches, len(num_blocks))
raise ValueError(error_msg)
 if num_branches != len(num_channels):
error_msg = 'NUM_BRANCHES({})<> NUM_CHANNELS({})'.format(
 num_branches, len(num_channels))
raise ValueError(error_msg)
 if num_branches != len(num_inchannels):
error_msg = 'NUM_BRANCHES({})<> NUM_INCHANNELS({})'.format(
 num_branches, len(num_inchannels))
raise ValueError(error_msg)
 def _make_one_branch(self, branch_index, block, num_blocks, num_channels,
stride=1):
downsample = None
if stride != 1 or \
self.num_inchannels[branch_index] != num_channels[branch_index] * block.expansion:
downsample = nn.Sequential(
 nn.Conv2d(self.num_inchannels[branch_index],
num_channels[branch_index] * block.expansion,
kernel_size=1, stride=stride, bias=False),
BatchNorm2d(num_channels[branch_index] * block.expansion,
momentum=BN_MOMENTUM), 
)
 layers = []
layers.append(block(self.num_inchannels[branch_index],
num_channels[branch_index], stride, downsample))
self.num_inchannels[branch_index] = \
num_channels[branch_index] * block.expansion
for i in range(1, num_blocks[branch_index]):
layers.append(block(self.num_inchannels[branch_index],
num_channels[branch_index]))
return nn.Sequential(*layers)
 def _make_branches(self, num_branches, block, num_blocks, num_channels):
branches = []
for i in range(num_branches):
branches.append(
 self._make_one_branch(i, block, num_blocks, num_channels))
return nn.ModuleList(branches)
 def _make_fuse_layers(self):
if self.num_branches == 1:
return None
num_branches = self.num_branches
num_inchannels = self.num_inchannels
fuse_layers = []
for i in range(num_branches if self.multi_scale_output else 1):
fuse_layer = []
for j in range(num_branches):
if j > i:
fuse_layer.append(nn.Sequential(
 nn.Conv2d(num_inchannels[j],
num_inchannels[i],
1,
1,
0,
bias=False),
BatchNorm2d(num_inchannels[i], momentum=BN_MOMENTUM)))
 elif j == i:
fuse_layer.append(None)
 else:
conv3x3s = []
for k in range(i-j):
if k == i - j - 1:
num_outchannels_conv3x3 = num_inchannels[i]
conv3x3s.append(nn.Sequential(
 nn.Conv2d(num_inchannels[j],
num_outchannels_conv3x3,
3, 2, 1, bias=False),
BatchNorm2d(num_outchannels_conv3x3,
momentum=BN_MOMENTUM)))
 else:
num_outchannels_conv3x3 = num_inchannels[j]
conv3x3s.append(nn.Sequential(
 nn.Conv2d(num_inchannels[j],
num_outchannels_conv3x3,
3, 2, 1, bias=False),
BatchNorm2d(num_outchannels_conv3x3,
momentum=BN_MOMENTUM),
nn.ReLU(inplace=relu_inplace)))
 fuse_layer.append(nn.Sequential(*conv3x3s))
fuse_layers.append(nn.ModuleList(fuse_layer))
return nn.ModuleList(fuse_layers)
 def get_num_inchannels(self):
return self.num_inchannels
def forward(self, x):
if self.num_branches == 1:
return [self.branches[0](x[0])]
for i in range(self.num_branches):
x[i] = self.branches[i](x[i])
 x_fuse = []
for i in range(len(self.fuse_layers)) :
y = x[0] if i == 0 else self.fuse_layers[i][0](x[0])
 for j in range(1, self.num_branches):
if i == j:
y = y + x[j]
elif j > i:
width_output = x[i].shape[-1]
height_output = x[i].shape[-2]
y = y + F.interpolate(
 self.fuse_layers[i][j](x[j]),
size=[height_output, width_output],
mode='bilinear', align_corners=ALIGN_CORNERS)
 else:
y = y + self.fuse_layers[i][j](x[j])
 x_fuse.append(self.relu(y))
return x_fuse
blocks_dict = {
'BASIC': BasicBlock,
'BOTTLENECK': Bottleneck
}
class ASPPModule(nn.Module):
def __init__(self, in_channel=720, hidden_channel=256, out_channel=512, dilations=(3, 6, 9), **kwargs):
super(ASPPModule, self ).__init__()
self.asp_conv2 = nn.Sequential(
 nn.Conv2d(in_channel, hidden_channel, kernel_size=1, padding=0, dilation=1, bias=True),
BatchNorm2d(hidden_channel),
nn.ReLU(inplace=relu_inplace),)
 self.asp_conv3 = nn.Sequential(
 nn.Conv2d(in_channel, hidden_channel, kernel_size=3, padding=dilations[0], dilation=dilations[0], bias=True),
BatchNorm2d(hidden_channel),
nn.ReLU(inplace=relu_inplace),)
 self.asp_conv4 = nn.Sequential(
 nn.Conv2d(in_channel, hidden_channel, kernel_size=3, padding=dilations[0], dilation=dilations[0], bias=True),
BatchNorm2d(hidden_channel),
nn.ReLU(inplace=relu_inplace),)
 self.asp_conv5 = nn.Sequential(
 nn.Conv2d(in_channel, hidden_channel, kernel_size=3, padding=dilations[0], dilation=dilations[0], bias=True),
BatchNorm2d(hidden_channel),
nn.ReLU(inplace=relu_inplace),)
 self.asp_fuse = nn.Sequential(
 nn.Conv2d(hidden_channel * 4 + 512, out_channel, kernel_size=1, padding=0, dilation=1, bias=True),
BatchNorm2d(out_channel),
nn.ReLU(inplace=relu_inplace),
nn.Dropout2d(0.1)
)
 def forward(self, feat_hr, feat_context):
feat1 = feat_context
feat2 = self.asp_conv2(feat_hr)
 feat3 = self.asp_conv3(feat_hr)
 feat4 = self.asp_conv4(feat_hr)
 feat5 = self.asp_conv5(feat_hr)
 out = torch.cat(( feat1, feat2, feat3, feat4, feat5), 1)
 output = self.asp_fuse(out)
 return output
class HighResolutionNet_ocr_aspp(nn.Module):
def __init__(self, config, dilations=(3, 6, 9), **kwargs):
global ALIGN_CORNERS
extra = config['MODEL']['EXTRA']
ocr_cfg = config['MODEL']['OCR']
super(HighResolutionNet_ocr_aspp, self ).__init__()
ALIGN_CORNERS = True
self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=2, padding=1,
bias=False)
 self.bn1 = BatchNorm2d(64, momentum=BN_MOMENTUM)
 self.conv2 = nn.Conv2d(64, 64, kernel_size=3, stride=2, padding=1,
bias=False)
 self.bn2 = BatchNorm2d(64, momentum=BN_MOMENTUM)
 self.relu = nn.ReLU(inplace=relu_inplace)
 self.stage1_cfg = extra['STAGE1']
num_channels = self.stage1_cfg['NUM_CHANNELS'][0]
block = blocks_dict[self.stage1_cfg['BLOCK']]
num_blocks = self.stage1_cfg['NUM_BLOCKS'][0]
self.layer1 = self._make_layer(block, 64, num_channels, num_blocks)
 stage1_out_channel = block.expansion * num_channels
self.stage2_cfg = extra['STAGE2']
num_channels = self.stage2_cfg['NUM_CHANNELS']
block = blocks_dict[self.stage2_cfg['BLOCK']]
num_channels = [
num_channels[i] * block.expansion for i in range(len(num_channels)) ]
self.transition1 = self._make_transition_layer(
 [stage1_out_channel], num_channels)
 self.stage2, pre_stage_channels = self._make_stage(
 self.stage2_cfg, num_channels)
 self.stage3_cfg = extra['STAGE3']
num_channels = self.stage3_cfg['NUM_CHANNELS']
block = blocks_dict[self.stage3_cfg['BLOCK']]
num_channels = [
num_channels[i] * block.expansion for i in range(len(num_channels)) ]
self.transition2 = self._make_transition_layer(
 pre_stage_channels, num_channels)
 self.stage3, pre_stage_channels = self._make_stage(
 self.stage3_cfg, num_channels)
 self.stage4_cfg = extra['STAGE4']
num_channels = self.stage4_cfg['NUM_CHANNELS']
block = blocks_dict[self.stage4_cfg['BLOCK']]
num_channels = [
num_channels[i] * block.expansion for i in range(len(num_channels)) ]
self.transition3 = self._make_transition_layer(
 pre_stage_channels, num_channels)
 self.stage4, pre_stage_channels = self._make_stage(
 self.stage4_cfg, num_channels, multi_scale_output=True)
 last_inp_channels = np.int(np.sum(pre_stage_channels))
ocr_mid_channels = ocr_cfg['MID_CHANNELS']
ocr_key_channels = ocr_cfg['KEY_CHANNELS']
self.conv3x3_ocr = nn.Sequential(
 nn.Conv2d(last_inp_channels, ocr_mid_channels,
kernel_size=3, stride=1, padding=1),
BatchNorm2d(ocr_mid_channels),
nn.ReLU(inplace=relu_inplace), 
)
 self.ocr_gather_head = SpatialGather_Module(config['DATASET']['NUM_CLASSES'])
 self.ocr_distri_head = SpatialOCR_Module(in_channels=ocr_mid_channels,
key_channels=ocr_key_channels,
out_channels=ocr_mid_channels,
scale=1,
dropout=0.05, 
)
 self.cls_head = nn.Conv2d(
 ocr_mid_channels, config['DATASET']['NUM_CLASSES'], kernel_size=1, stride=1, padding=0, bias=True)
 self.aux_head = nn.Sequential(
 nn.Conv2d(last_inp_channels, last_inp_channels,
kernel_size=1, stride=1, padding=0),
BatchNorm2d(last_inp_channels),
nn.ReLU(inplace=relu_inplace),
nn.Conv2d(last_inp_channels, config['DATASET']['NUM_CLASSES'],
kernel_size=1, stride=1, padding=0, bias=True)
)
 self.asp_module = ASPPModule(in_channel=720, hidden_channel=256, out_channel=512, dilations=dilations)
 self.backbone = nn.ModuleList([self.conv1, self.bn1, self.conv2, self.bn2, self.layer1, self.transition1,
self.stage2, self.transition2, self.stage3, self.transition3, self.stage4])
 self.fuse = nn.ModuleList(
 [self.aux_head, self.conv3x3_ocr, self.ocr_gather_head, self.ocr_distri_head, self.cls_head,
self.asp_module])
 def _make_transition_layer(
 self, num_channels_pre_layer, num_channels_cur_layer):
num_branches_cur = len(num_channels_cur_layer)
 num_branches_pre = len(num_channels_pre_layer)
 transition_layers = []
for i in range(num_branches_cur):
if i < num_branches_pre:
if num_channels_cur_layer[i] != num_channels_pre_layer[i]:
transition_layers.append(nn.Sequential(
 nn.Conv2d(num_channels_pre_layer[i],
num_channels_cur_layer[i],
3,
1,
1,
bias=False),
BatchNorm2d(
 num_channels_cur_layer[i], momentum=BN_MOMENTUM),
nn.ReLU(inplace=relu_inplace)))
 else:
transition_layers.append(None)
 else:
conv3x3s = []
for j in range(i + 1 - num_branches_pre):
inchannels = num_channels_pre_layer[-1]
outchannels = num_channels_cur_layer[i] \
if j == i - num_branches_pre else inchannels
conv3x3s.append(nn.Sequential(
 nn.Conv2d(
 inchannels, outchannels, 3, 2, 1, bias=False),
BatchNorm2d(outchannels, momentum=BN_MOMENTUM),
nn.ReLU(inplace=relu_inplace)))
 transition_layers.append(nn.Sequential(*conv3x3s))
return nn.ModuleList(transition_layers)
 def _make_layer(self, block, inplanes, planes, blocks, stride=1):
downsample = None
if stride != 1 or inplanes != planes * block.expansion:
downsample = nn.Sequential(
 nn.Conv2d(inplanes, planes * block.expansion,
kernel_size=1, stride=stride, bias=False),
BatchNorm2d(planes * block.expansion, momentum=BN_MOMENTUM), 
)
 layers = []
layers.append(block(inplanes, planes, stride, downsample))
inplanes = planes * block.expansion
for i in range(1, blocks):
layers.append(block(inplanes, planes))
return nn.Sequential(*layers)
 def _make_stage(self, layer_config, num_inchannels,
multi_scale_output=True):
num_modules = layer_config['NUM_MODULES']
num_branches = layer_config['NUM_BRANCHES']
num_blocks = layer_config['NUM_BLOCKS']
num_channels = layer_config['NUM_CHANNELS']
block = blocks_dict[layer_config['BLOCK']]
fuse_method = layer_config['FUSE_METHOD']
modules = []
for i in range(num_modules):
if not multi_scale_output and i == num_modules - 1:
reset_multi_scale_output = False
else:
reset_multi_scale_output = True
modules.append(
 HighResolutionModule(num_branches,
block,
num_blocks,
num_inchannels,
num_channels,
fuse_method,
reset_multi_scale_output)
)
 num_inchannels = modules[-1].get_num_inchannels()
return nn.Sequential(*modules), num_inchannels
def forward(self, x):
x = self.conv1(x)
 x = self.bn1(x)
 x = self.relu(x)
 x = self.conv2(x)
 x = self.bn2(x)
 x = self.relu(x)
 x = self.layer1(x)
 x_list = []
for i in range(self.stage2_cfg['NUM_BRANCHES']):
if self.transition1[i] is not None:
x_list.append(self.transition1[i](x))
else:
x_list.append(x)
 y_list = self.stage2(x_list)
 x_list = []
for i in range(self.stage3_cfg['NUM_BRANCHES']):
if self.transition2[i] is not None:
if i < self.stage2_cfg['NUM_BRANCHES']:
x_list.append(self.transition2[i](y_list[i]))
else:
x_list.append(self.transition2[i](y_list[-1]))
else:
x_list.append(y_list[i])
 y_list = self.stage3(x_list)
 x_list = []
for i in range(self.stage4_cfg['NUM_BRANCHES']):
if self.transition3[i] is not None:
if i < self.stage3_cfg['NUM_BRANCHES']:
x_list.append(self.transition3[i](y_list[i]))
else:
x_list.append(self.transition3[i](y_list[-1]))
else:
x_list.append(y_list[i])
 x = self.stage4(x_list)
 x0_h, x0_w = x[0].size(2), x[0].size(3)
 x1 = F.interpolate(x[1], size=(x0_h, x0_w),
mode='bilinear', align_corners=ALIGN_CORNERS)
 x2 = F.interpolate(x[2], size=(x0_h, x0_w),
mode='bilinear', align_corners=ALIGN_CORNERS)
 x3 = F.interpolate(x[3], size=(x0_h, x0_w),
mode='bilinear', align_corners=ALIGN_CORNERS)
 feats_hr = torch.cat([x[0], x1, x2, x3], 1)
 out_aux_seg = []
out_aux = self.aux_head(feats_hr)
 feats = self.conv3x3_ocr(feats_hr)
 context = self.ocr_gather_head(feats, out_aux)
 feats_context = self.ocr_distri_head(feats, context)
 feats = self.asp_module(feats_hr, feats_context)
 out = self.cls_head(feats)
 out_aux_seg.append(out_aux)
 out_aux_seg.append(out)
 return out_aux_seg
def init_weights(self, pretrained='',):
print('=> init weights from normal distribution')
 for name, m in self.named_modules() :
if any(part in name for part in {'cls', 'aux', 'ocr'}):
continue
if isinstance(m, nn.Conv2d):
nn.init.normal_(m.weight, std=0.001)
 elif isinstance(m, BatchNorm2d_class):
nn.init.constant_(m.weight, 1)
 nn.init.constant_(m.bias, 0)
 if os.path.isfile(pretrained):
pretrained_dict = torch.load(pretrained, map_location={'cuda:0': 'cpu'})
 print('=> loading pretrained model {}'.format(pretrained))
model_dict = self.state_dict()
pretrained_dict = {k.replace('last_layer', 'aux_head' ).replace('model.', ''): v for k, v in
pretrained_dict.items() }
pretrained_dict = {k: v for k, v in pretrained_dict.items()
if k in model_dict.keys() }
print('length of pretrained_dict:', len(pretrained_dict.items()))
model_dict.update(pretrained_dict)
 self.load_state_dict(model_dict)
 elif pretrained:
raise RuntimeError('No such file {}'.format(pretrained))
def trainable_parameters(self):
return list(self.backbone.parameters()), list(self.fuse.parameters())
 def get_hr_ocr_asp_seg_model(cfg, **kwargs):
cfg['MODEL']['NUM_CLASSES'] = kwargs['num_classes']
model = HighResolutionNet_ocr_aspp(cfg, **kwargs)
 return model
def val_epoch(model, val_iterator, criterion, device="cpu"):
model.eval()
train_loss = []
num_corrects = 0
num_total = 0
labels = []
outs = []
tbar = tqdm(val_iterator)
 for item in tbar:
x = item[0].to(device ).long()
target_id = item[1].to(device ).long()
label = item[2].to(device ).float()
target_mask =(target_id != 0)
 with torch.no_grad() :
output, atten_weight = model(x, target_id)
 output = torch.masked_select(output, target_mask)
 label = torch.masked_select(label, target_mask)
 loss = criterion(output, label)
 train_loss.append(loss.item())
 pred =(torch.sigmoid(output)>= 0.5 ).long()
num_corrects +=(pred == label ).sum().item()
num_total += len(label)
 labels.extend(label.view(-1 ).data.cpu().numpy())
 outs.extend(output.view(-1 ).data.cpu().numpy())
 tbar.set_description('loss - {:.4f}'.format(loss))
acc = num_corrects / num_total
auc = roc_auc_score(labels, outs)
 loss = np.average(train_loss)
 return loss, acc, auc<train_model>
def INF(B, H, W):
return -torch.diag(torch.tensor(float("inf")).cuda().repeat(H), 0 ).unsqueeze(0 ).repeat(B * W, 1, 1)
 class CC_module(nn.Module):
def __init__(self, in_dim):
super(CC_module, self ).__init__()
self.query_conv = nn.Conv2d(in_channels=in_dim, out_channels=in_dim // 8, kernel_size=1)
 self.key_conv = nn.Conv2d(in_channels=in_dim, out_channels=in_dim // 8, kernel_size=1)
 self.value_conv = nn.Conv2d(in_channels=in_dim, out_channels=in_dim, kernel_size=1)
 self.softmax = Softmax(dim=3)
 self.INF = INF
self.gamma = nn.Parameter(torch.zeros(1))
def forward(self, x):
m_batchsize, _, height, width = x.size()
proj_query = self.query_conv(x)
 proj_query_H = proj_query.permute(0, 3, 1, 2 ).contiguous().view(m_batchsize * width, -1, height ).permute(0, 2,
1)
 proj_query_W = proj_query.permute(0, 2, 1, 3 ).contiguous().view(m_batchsize * height, -1, width ).permute(0, 2,
1)
 proj_key = self.key_conv(x)
 proj_key_H = proj_key.permute(0, 3, 1, 2 ).contiguous().view(m_batchsize * width, -1, height)
 proj_key_W = proj_key.permute(0, 2, 1, 3 ).contiguous().view(m_batchsize * height, -1, width)
 proj_value = self.value_conv(x)
 proj_value_H = proj_value.permute(0, 3, 1, 2 ).contiguous().view(m_batchsize * width, -1, height)
 proj_value_W = proj_value.permute(0, 2, 1, 3 ).contiguous().view(m_batchsize * height, -1, width)
 energy_H =(torch.bmm(proj_query_H, proj_key_H)+ self.INF(m_batchsize, height, width)).view(m_batchsize, width,
height,
height ).permute(0,
2,
1,
3)
 energy_W = torch.bmm(proj_query_W, proj_key_W ).view(m_batchsize, height, width, width)
 concate = self.softmax(torch.cat([energy_H, energy_W], 3))
att_H = concate[:, :, :, 0:height].permute(0, 2, 1, 3 ).contiguous().view(m_batchsize * width, height, height)
 att_W = concate[:, :, :, height:height + width].contiguous().view(m_batchsize * height, width, width)
 out_H = torch.bmm(proj_value_H, att_H.permute(0, 2, 1)).view(m_batchsize, width, -1, height ).permute(0, 2, 3, 1)
 out_W = torch.bmm(proj_value_W, att_W.permute(0, 2, 1)).view(m_batchsize, height, -1, width ).permute(0, 2, 1, 3)
 return self.gamma *(out_H + out_W)+ x
class RCCAModule(nn.Module):
def __init__(self, in_channels, out_channels, num_classes):
super(RCCAModule, self ).__init__()
inter_channels = in_channels // 4
self.conva = nn.Sequential(nn.Conv2d(in_channels, inter_channels, 3, padding=1, bias=False),
nn.BatchNorm2d(inter_channels))
self.cca = CC_module(inter_channels)
 self.convb = nn.Sequential(nn.Conv2d(inter_channels, inter_channels, 3, padding=1, bias=False),
nn.BatchNorm2d(inter_channels))
self.bottleneck = nn.Sequential(
 nn.Conv2d(in_channels+inter_channels, out_channels, kernel_size=3, padding=1, dilation=1, bias=False),
nn.BatchNorm2d(out_channels),
nn.Dropout2d(0.1),
nn.Conv2d(512, num_classes, kernel_size=1, stride=1, padding=0, bias=True)
)
 def forward(self, x, recurrence=1):
output = self.conva(x)
 for i in range(recurrence):
output = self.cca(output)
 output = self.convb(output)
 output = self.bottleneck(torch.cat([x, output], 1))
return output
class BasicBlock(nn.Module):
expansion = 1
def __init__(self, inplanes, planes, stride=1, downsample=None):
super(BasicBlock, self ).__init__()
self.conv1 = conv3x3(inplanes, planes, stride)
 self.bn1 = BatchNorm2d(planes, momentum=BN_MOMENTUM)
 self.relu = nn.ReLU(inplace=relu_inplace)
 self.conv2 = conv3x3(planes, planes)
 self.bn2 = BatchNorm2d(planes, momentum=BN_MOMENTUM)
 self.downsample = downsample
self.stride = stride
def forward(self, x):
residual = x
out = self.conv1(x)
 out = self.bn1(out)
 out = self.relu(out)
 out = self.conv2(out)
 out = self.bn2(out)
 if self.downsample is not None:
residual = self.downsample(x)
 out = out + residual
out = self.relu(out)
 return out
class Bottleneck(nn.Module):
expansion = 4
def __init__(self, inplanes, planes, stride=1, downsample=None):
super(Bottleneck, self ).__init__()
self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
 self.bn1 = BatchNorm2d(planes, momentum=BN_MOMENTUM)
 self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
padding=1, bias=False)
 self.bn2 = BatchNorm2d(planes, momentum=BN_MOMENTUM)
 self.conv3 = nn.Conv2d(planes, planes * self.expansion, kernel_size=1,
bias=False)
 self.bn3 = BatchNorm2d(planes * self.expansion,
momentum=BN_MOMENTUM)
 self.relu = nn.ReLU(inplace=relu_inplace)
 self.downsample = downsample
self.stride = stride
def forward(self, x):
residual = x
out = self.conv1(x)
 out = self.bn1(out)
 out = self.relu(out)
 out = self.conv2(out)
 out = self.bn2(out)
 out = self.relu(out)
 out = self.conv3(out)
 out = self.bn3(out)
 if self.downsample is not None:
residual = self.downsample(x)
 out = out + residual
out = self.relu(out)
 return out
class HighResolutionModule(nn.Module):
def __init__(self, num_branches, blocks, num_blocks, num_inchannels,
num_channels, fuse_method, multi_scale_output=True):
super(HighResolutionModule, self ).__init__()
self._check_branches(
 num_branches, blocks, num_blocks, num_inchannels, num_channels)
 self.num_inchannels = num_inchannels
self.fuse_method = fuse_method
self.num_branches = num_branches
self.multi_scale_output = multi_scale_output
self.branches = self._make_branches(
 num_branches, blocks, num_blocks, num_channels)
 self.fuse_layers = self._make_fuse_layers()
self.relu = nn.ReLU(inplace=relu_inplace)
 def _check_branches(self, num_branches, blocks, num_blocks,
num_inchannels, num_channels):
if num_branches != len(num_blocks):
error_msg = 'NUM_BRANCHES({})<> NUM_BLOCKS({})'.format(
 num_branches, len(num_blocks))
raise ValueError(error_msg)
 if num_branches != len(num_channels):
error_msg = 'NUM_BRANCHES({})<> NUM_CHANNELS({})'.format(
 num_branches, len(num_channels))
raise ValueError(error_msg)
 if num_branches != len(num_inchannels):
error_msg = 'NUM_BRANCHES({})<> NUM_INCHANNELS({})'.format(
 num_branches, len(num_inchannels))
raise ValueError(error_msg)
 def _make_one_branch(self, branch_index, block, num_blocks, num_channels,
stride=1):
downsample = None
if stride != 1 or \
self.num_inchannels[branch_index] != num_channels[branch_index] * block.expansion:
downsample = nn.Sequential(
 nn.Conv2d(self.num_inchannels[branch_index],
num_channels[branch_index] * block.expansion,
kernel_size=1, stride=stride, bias=False),
BatchNorm2d(num_channels[branch_index] * block.expansion,
momentum=BN_MOMENTUM), 
)
 layers = []
layers.append(block(self.num_inchannels[branch_index],
num_channels[branch_index], stride, downsample))
self.num_inchannels[branch_index] = \
num_channels[branch_index] * block.expansion
for i in range(1, num_blocks[branch_index]):
layers.append(block(self.num_inchannels[branch_index],
num_channels[branch_index]))
return nn.Sequential(*layers)
 def _make_branches(self, num_branches, block, num_blocks, num_channels):
branches = []
for i in range(num_branches):
branches.append(
 self._make_one_branch(i, block, num_blocks, num_channels))
return nn.ModuleList(branches)
 def _make_fuse_layers(self):
if self.num_branches == 1:
return None
num_branches = self.num_branches
num_inchannels = self.num_inchannels
fuse_layers = []
for i in range(num_branches if self.multi_scale_output else 1):
fuse_layer = []
for j in range(num_branches):
if j > i:
fuse_layer.append(nn.Sequential(
 nn.Conv2d(num_inchannels[j],
num_inchannels[i],
1,
1,
0,
bias=False),
BatchNorm2d(num_inchannels[i], momentum=BN_MOMENTUM)))
 elif j == i:
fuse_layer.append(None)
 else:
conv3x3s = []
for k in range(i - j):
if k == i - j - 1:
num_outchannels_conv3x3 = num_inchannels[i]
conv3x3s.append(nn.Sequential(
 nn.Conv2d(num_inchannels[j],
num_outchannels_conv3x3,
3, 2, 1, bias=False),
BatchNorm2d(num_outchannels_conv3x3,
momentum=BN_MOMENTUM)))
 else:
num_outchannels_conv3x3 = num_inchannels[j]
conv3x3s.append(nn.Sequential(
 nn.Conv2d(num_inchannels[j],
num_outchannels_conv3x3,
3, 2, 1, bias=False),
BatchNorm2d(num_outchannels_conv3x3,
momentum=BN_MOMENTUM),
nn.ReLU(inplace=relu_inplace)))
 fuse_layer.append(nn.Sequential(*conv3x3s))
fuse_layers.append(nn.ModuleList(fuse_layer))
return nn.ModuleList(fuse_layers)
 def get_num_inchannels(self):
return self.num_inchannels
def forward(self, x):
if self.num_branches == 1:
return [self.branches[0](x[0])]
for i in range(self.num_branches):
x[i] = self.branches[i](x[i])
 x_fuse = []
for i in range(len(self.fuse_layers)) :
y = x[0] if i == 0 else self.fuse_layers[i][0](x[0])
 for j in range(1, self.num_branches):
if i == j:
y = y + x[j]
elif j > i:
width_output = x[i].shape[-1]
height_output = x[i].shape[-2]
y = y + F.interpolate(
 self.fuse_layers[i][j](x[j]),
size=[height_output, width_output],
mode='bilinear', align_corners=ALIGN_CORNERS)
 else:
y = y + self.fuse_layers[i][j](x[j])
 x_fuse.append(self.relu(y))
return x_fuse
blocks_dict = {
'BASIC': BasicBlock,
'BOTTLENECK': Bottleneck
}
class HighResolutionNetCC(nn.Module):
def __init__(self, config, **kwargs):
global ALIGN_CORNERS
extra = config['MODEL']['EXTRA']
ocr_cfg = config['MODEL']['OCR']
super(HighResolutionNetCC, self ).__init__()
ALIGN_CORNERS = True
self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=2, padding=1,
bias=False)
 self.bn1 = BatchNorm2d(64, momentum=BN_MOMENTUM)
 self.conv2 = nn.Conv2d(64, 64, kernel_size=3, stride=2, padding=1,
bias=False)
 self.bn2 = BatchNorm2d(64, momentum=BN_MOMENTUM)
 self.relu = nn.ReLU(inplace=relu_inplace)
 self.stage1_cfg = extra['STAGE1']
num_channels = self.stage1_cfg['NUM_CHANNELS'][0]
block = blocks_dict[self.stage1_cfg['BLOCK']]
num_blocks = self.stage1_cfg['NUM_BLOCKS'][0]
self.layer1 = self._make_layer(block, 64, num_channels, num_blocks)
 stage1_out_channel = block.expansion * num_channels
self.stage2_cfg = extra['STAGE2']
num_channels = self.stage2_cfg['NUM_CHANNELS']
block = blocks_dict[self.stage2_cfg['BLOCK']]
num_channels = [
num_channels[i] * block.expansion for i in range(len(num_channels)) ]
self.transition1 = self._make_transition_layer(
 [stage1_out_channel], num_channels)
 self.stage2, pre_stage_channels = self._make_stage(
 self.stage2_cfg, num_channels)
 self.stage3_cfg = extra['STAGE3']
num_channels = self.stage3_cfg['NUM_CHANNELS']
block = blocks_dict[self.stage3_cfg['BLOCK']]
num_channels = [
num_channels[i] * block.expansion for i in range(len(num_channels)) ]
self.transition2 = self._make_transition_layer(
 pre_stage_channels, num_channels)
 self.stage3, pre_stage_channels = self._make_stage(
 self.stage3_cfg, num_channels)
 self.stage4_cfg = extra['STAGE4']
num_channels = self.stage4_cfg['NUM_CHANNELS']
block = blocks_dict[self.stage4_cfg['BLOCK']]
num_channels = [
num_channels[i] * block.expansion for i in range(len(num_channels)) ]
self.transition3 = self._make_transition_layer(
 pre_stage_channels, num_channels)
 self.stage4, pre_stage_channels = self._make_stage(
 self.stage4_cfg, num_channels, multi_scale_output=True)
 last_inp_channels = np.int(np.sum(pre_stage_channels))
ocr_mid_channels = ocr_cfg['MID_CHANNELS']
ocr_key_channels = ocr_cfg['KEY_CHANNELS']
self.conv3x3_ocr = nn.Sequential(
 nn.Conv2d(last_inp_channels, ocr_mid_channels,
kernel_size=3, stride=1, padding=1),
BatchNorm2d(ocr_mid_channels),
nn.ReLU(inplace=relu_inplace), 
)
 self.cc = RCCAModule(512, 512, 5)
 self.aux_head = nn.Sequential(
 nn.Conv2d(last_inp_channels, last_inp_channels,
kernel_size=1, stride=1, padding=0),
BatchNorm2d(last_inp_channels),
nn.ReLU(inplace=relu_inplace),
nn.Conv2d(last_inp_channels, config['DATASET']['NUM_CLASSES'],
kernel_size=1, stride=1, padding=0, bias=True)
)
 self.backbone = nn.ModuleList([self.conv1, self.bn1, self.conv2, self.bn2, self.layer1, self.transition1,
self.stage2, self.transition2, self.stage3, self.transition3, self.stage4])
 self.fuse = nn.ModuleList(
 [self.aux_head, self.conv3x3_ocr, self.cc])
 def _make_transition_layer(
 self, num_channels_pre_layer, num_channels_cur_layer):
num_branches_cur = len(num_channels_cur_layer)
 num_branches_pre = len(num_channels_pre_layer)
 transition_layers = []
for i in range(num_branches_cur):
if i < num_branches_pre:
if num_channels_cur_layer[i] != num_channels_pre_layer[i]:
transition_layers.append(nn.Sequential(
 nn.Conv2d(num_channels_pre_layer[i],
num_channels_cur_layer[i],
3,
1,
1,
bias=False),
BatchNorm2d(
 num_channels_cur_layer[i], momentum=BN_MOMENTUM),
nn.ReLU(inplace=relu_inplace)))
 else:
transition_layers.append(None)
 else:
conv3x3s = []
for j in range(i + 1 - num_branches_pre):
inchannels = num_channels_pre_layer[-1]
outchannels = num_channels_cur_layer[i] \
if j == i - num_branches_pre else inchannels
conv3x3s.append(nn.Sequential(
 nn.Conv2d(
 inchannels, outchannels, 3, 2, 1, bias=False),
BatchNorm2d(outchannels, momentum=BN_MOMENTUM),
nn.ReLU(inplace=relu_inplace)))
 transition_layers.append(nn.Sequential(*conv3x3s))
return nn.ModuleList(transition_layers)
 def _make_layer(self, block, inplanes, planes, blocks, stride=1):
downsample = None
if stride != 1 or inplanes != planes * block.expansion:
downsample = nn.Sequential(
 nn.Conv2d(inplanes, planes * block.expansion,
kernel_size=1, stride=stride, bias=False),
BatchNorm2d(planes * block.expansion, momentum=BN_MOMENTUM), 
)
 layers = []
layers.append(block(inplanes, planes, stride, downsample))
inplanes = planes * block.expansion
for i in range(1, blocks):
layers.append(block(inplanes, planes))
return nn.Sequential(*layers)
 def _make_stage(self, layer_config, num_inchannels,
multi_scale_output=True):
num_modules = layer_config['NUM_MODULES']
num_branches = layer_config['NUM_BRANCHES']
num_blocks = layer_config['NUM_BLOCKS']
num_channels = layer_config['NUM_CHANNELS']
block = blocks_dict[layer_config['BLOCK']]
fuse_method = layer_config['FUSE_METHOD']
modules = []
for i in range(num_modules):
if not multi_scale_output and i == num_modules - 1:
reset_multi_scale_output = False
else:
reset_multi_scale_output = True
modules.append(
 HighResolutionModule(num_branches,
block,
num_blocks,
num_inchannels,
num_channels,
fuse_method,
reset_multi_scale_output)
)
 num_inchannels = modules[-1].get_num_inchannels()
return nn.Sequential(*modules), num_inchannels
def forward(self, x):
x = self.conv1(x)
 x = self.bn1(x)
 x = self.relu(x)
 x = self.conv2(x)
 x = self.bn2(x)
 x = self.relu(x)
 x = self.layer1(x)
 x_list = []
for i in range(self.stage2_cfg['NUM_BRANCHES']):
if self.transition1[i] is not None:
x_list.append(self.transition1[i](x))
else:
x_list.append(x)
 y_list = self.stage2(x_list)
 x_list = []
for i in range(self.stage3_cfg['NUM_BRANCHES']):
if self.transition2[i] is not None:
if i < self.stage2_cfg['NUM_BRANCHES']:
x_list.append(self.transition2[i](y_list[i]))
else:
x_list.append(self.transition2[i](y_list[-1]))
else:
x_list.append(y_list[i])
 y_list = self.stage3(x_list)
 x_list = []
for i in range(self.stage4_cfg['NUM_BRANCHES']):
if self.transition3[i] is not None:
if i < self.stage3_cfg['NUM_BRANCHES']:
x_list.append(self.transition3[i](y_list[i]))
else:
x_list.append(self.transition3[i](y_list[-1]))
else:
x_list.append(y_list[i])
 x = self.stage4(x_list)
 x0_h, x0_w = x[0].size(2), x[0].size(3)
 x1 = F.interpolate(x[1], size=(x0_h, x0_w),
mode='bilinear', align_corners=ALIGN_CORNERS)
 x2 = F.interpolate(x[2], size=(x0_h, x0_w),
mode='bilinear', align_corners=ALIGN_CORNERS)
 x3 = F.interpolate(x[3], size=(x0_h, x0_w),
mode='bilinear', align_corners=ALIGN_CORNERS)
 feats = torch.cat([x[0], x1, x2, x3], 1)
 out_aux_seg = []
out_aux = self.aux_head(feats)
 feats = self.conv3x3_ocr(feats)
 out = self.cc(feats)
 out_aux_seg.append(out_aux)
 out_aux_seg.append(out)
 return out_aux_seg
def init_weights(self, pretrained='',):
print('=> init weights from normal distribution')
 for name, m in self.named_modules() :
if any(part in name for part in {'cls', 'aux', 'ocr'}):
continue
if isinstance(m, nn.Conv2d):
nn.init.normal_(m.weight, std=0.001)
 elif isinstance(m, BatchNorm2d_class):
nn.init.constant_(m.weight, 1)
 nn.init.constant_(m.bias, 0)
 if os.path.isfile(pretrained):
pretrained_dict = torch.load(pretrained, map_location={'cuda:0': 'cpu'})
 print('=> loading pretrained model {}'.format(pretrained))
model_dict = self.state_dict()
pretrained_dict = {k.replace('last_layer', 'aux_head' ).replace('model.', ''): v for k, v in
pretrained_dict.items() }
pretrained_dict = {k: v for k, v in pretrained_dict.items()
if k in model_dict.keys() }
print('length of pretrained_dict:', len(pretrained_dict.items()))
model_dict.update(pretrained_dict)
 self.load_state_dict(model_dict)
 elif pretrained:
raise RuntimeError('No such file {}'.format(pretrained))
def trainable_parameters(self):
return list(self.backbone.parameters()), list(self.fuse.parameters())
 def get_seg_cc_model(cfg, **kwargs):
cfg['MODEL']['NUM_CLASSES'] = kwargs['num_classes']
model = HighResolutionNetCC(cfg, **kwargs)
 return model
epochs = 20
over_fit = 0
last_auc = 0
for epoch in range(epochs):
train_loss, train_acc, train_auc = train_epoch(model, train_dataloader, optimizer, criterion, device)
 print("epoch - {} train_loss - {:.2f} acc - {:.3f} auc - {:.3f}".format(epoch, train_loss, train_acc, train_auc))
val_loss, avl_acc, val_auc = val_epoch(model, val_dataloader, criterion, device)
 print("epoch - {} val_loss - {:.2f} acc - {:.3f} auc - {:.3f}".format(epoch, val_loss, avl_acc, val_auc))
if val_auc > last_auc:
last_auc = val_auc
over_fit = 0
else:
over_fit += 1
if over_fit >= 2:
print("early stop epoch ", epoch)
 break<split>
class Bottle2neck(nn.Module):
expansion = 4
def __init__(self, inplanes, planes, stride=1, downsample=None, baseWidth=26, scale=4, stype='normal'):

 super(Bottle2neck, self ).__init__()
width = int(math.floor(planes *(baseWidth / 64.0)))
 self.conv1 = nn.Conv2d(inplanes, width * scale, kernel_size=1, bias=False)
 self.bn1 = nn.BatchNorm2d(width * scale)
 if scale == 1:
self.nums = 1
else:
self.nums = scale - 1
if stype == 'stage':
self.pool = nn.AvgPool2d(kernel_size=3, stride=stride, padding=1)
 convs = []
bns = []
for i in range(self.nums):
convs.append(nn.Conv2d(width, width, kernel_size=3, stride=stride, padding=1, bias=False))
bns.append(nn.BatchNorm2d(width))
self.convs = nn.ModuleList(convs)
 self.bns = nn.ModuleList(bns)
 self.conv3 = nn.Conv2d(width * scale, planes * self.expansion, kernel_size=1, bias=False)
 self.bn3 = nn.BatchNorm2d(planes * self.expansion)
 self.relu = nn.ReLU(inplace=True)
 self.downsample = downsample
self.stype = stype
self.scale = scale
self.width = width
def forward(self, x):
residual = x
out = self.conv1(x)
 out = self.bn1(out)
 out = self.relu(out)
 spx = torch.split(out, self.width, 1)
 for i, _ in enumerate(range(self.nums)) :
if i == 0 or self.stype == 'stage':
sp = spx[i]
else:
sp = sp + spx[i]
sp = self.convs[i](sp)
 sp = self.relu(self.bns[i](sp))
if i == 0:
out = sp
else:
out = torch.cat(( out, sp), 1)
 if self.scale != 1 and self.stype == 'normal':
out = torch.cat(( out, spx[self.nums]), 1)
 elif self.scale != 1 and self.stype == 'stage':
out = torch.cat(( out, self.pool(spx[self.nums])) , 1)
 out = self.conv3(out)
 out = self.bn3(out)
 if self.downsample is not None:
residual = self.downsample(x)
 out += residual
out = self.relu(out)
 return out
class Res2Net(nn.Module):
def __init__(self, block, layers, baseWidth=26, scale=4, num_classes=1000):
self.inplanes = 64
super(Res2Net, self ).__init__()
self.baseWidth = baseWidth
self.scale = scale
self.conv1 = nn.Sequential(
 nn.Conv2d(3, 32, 3, 2, 1, bias=False),
nn.BatchNorm2d(32),
nn.ReLU(inplace=True),
nn.Conv2d(32, 32, 3, 1, 1, bias=False),
nn.BatchNorm2d(32),
nn.ReLU(inplace=True),
nn.Conv2d(32, 64, 3, 1, 1, bias=False)
)
 self.bn1 = nn.BatchNorm2d(64)
 self.relu = nn.ReLU()
self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
 self.layer1 = self._make_layer(block, 64, layers[0])
 self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
 self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
 self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
 self.avgpool = nn.AdaptiveAvgPool2d(1)
 self.fc = nn.Linear(512 * block.expansion, num_classes)
 for m in self.modules() :
if isinstance(m, nn.Conv2d):
nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
 elif isinstance(m, nn.BatchNorm2d):
nn.init.constant_(m.weight, 1)
 nn.init.constant_(m.bias, 0)
 def _make_layer(self, block, planes, blocks, stride=1):
downsample = None
if stride != 1 or self.inplanes != planes * block.expansion:
downsample = nn.Sequential(
 nn.AvgPool2d(kernel_size=stride, stride=stride,
ceil_mode=True, count_include_pad=False),
nn.Conv2d(self.inplanes, planes * block.expansion,
kernel_size=1, stride=1, bias=False),
nn.BatchNorm2d(planes * block.expansion), 
)
 layers = []
layers.append(block(self.inplanes, planes, stride, downsample=downsample,
stype='stage', baseWidth=self.baseWidth, scale=self.scale))
self.inplanes = planes * block.expansion
for i in range(1, blocks):
layers.append(block(self.inplanes, planes, baseWidth=self.baseWidth, scale=self.scale))
return nn.Sequential(*layers)
 def forward(self, x):
x = self.conv1(x)
 x = self.bn1(x)
 x = self.relu(x)
 x = self.maxpool(x)
 x = self.layer1(x)
 x = self.layer2(x)
 x = self.layer3(x)
 x = self.layer4(x)
 x = self.avgpool(x)
 x = x.view(x.size(0), -1)
 x = self.fc(x)
 return x
def res2net101_v1b(pretrained=False, **kwargs):

 model = Res2Net(Bottle2neck, [3, 4, 23, 3], baseWidth=26, scale=4, **kwargs)
 return model
env = riiideducation.make_env()
iter_test = env.iter_test()<feature_engineering>
class Res2NetDeeplabV3P(nn.Module):
def __init__(self, num_classes, aspp_dilate=(12, 24, 36)) :
super(Res2NetDeeplabV3P, self ).__init__()
backbone = res2net101_v1b(True)
 self.conv1 = backbone.conv1
self.bn1 = backbone.bn1
self.relu = backbone.relu
self.maxpool = backbone.maxpool
self.stage1 = backbone.layer1
self.stage2 = backbone.layer2
self.stage3 = backbone.layer3
self.stage4 = backbone.layer4
self.project = nn.Sequential(
 nn.Conv2d(512, 48, 1, bias=False),
nn.BatchNorm2d(48),
nn.ReLU(inplace=True), 
)
 self.aspp = ASPP(2048, aspp_dilate)
 self.classifier = nn.Sequential(
 nn.Conv2d(304, 256, 3, padding=1, bias=False),
nn.BatchNorm2d(256),
nn.ReLU(inplace=True),
nn.Conv2d(256, num_classes, 1)
)
 self.backbone = nn.ModuleList([self.conv1, self.bn1, self.stage1, self.stage2, self.stage3, self.stage4])
 self.fuse = nn.ModuleList([self.project, self.aspp, self.classifier])
 self._init_fuse_weight()
def trainable_parameters(self):
return list(self.backbone.parameters()), list(self.fuse.parameters())
 def forward(self, x):
x = self.conv1(x)
 x = self.bn1(x)
 x = self.relu(x)
 x = self.maxpool(x)
 x = self.stage1(x)
 x = self.stage2(x)
 low_feat = x.clone()
x = self.stage3(x)
 out_feat = self.stage4(x)
 low_level_feature = self.project(low_feat)
 output_feature = self.aspp(out_feat)
 output_feature = F.interpolate(output_feature, size=low_level_feature.shape[2:], mode='bilinear',
align_corners=False)
 return self.classifier(torch.cat([low_level_feature, output_feature], dim=1))
def _init_fuse_weight(self):
for m in self.fuse.modules() :
if isinstance(m, nn.Conv2d):
nn.init.kaiming_normal_(m.weight)
 elif isinstance(m,(nn.BatchNorm2d, nn.GroupNorm)) :
nn.init.constant_(m.weight, 1)
 nn.init.constant_(m.bias, 0)
 class ASPPConv(nn.Sequential):
def __init__(self, in_channels, out_channels, dilation):
modules = [
nn.Conv2d(in_channels, out_channels, 3, padding=dilation, dilation=dilation, bias=False),
nn.BatchNorm2d(out_channels),
nn.ReLU(inplace=True)
 ]
super(ASPPConv, self ).__init__(*modules)
 class ASPPPooling(nn.Sequential):
def __init__(self, in_channels, out_channels):
super(ASPPPooling, self ).__init__(
 nn.AdaptiveAvgPool2d(1),
nn.Conv2d(in_channels, out_channels, 1, bias=False),
nn.BatchNorm2d(out_channels),
nn.ReLU(inplace=True))
def forward(self, x):
size = x.shape[-2:]
x = super(ASPPPooling, self ).forward(x)
 return F.interpolate(x, size=size, mode='bilinear', align_corners=False)
 class ASPP(nn.Module):
def __init__(self, in_channels, atrous_rates):
super(ASPP, self ).__init__()
out_channels = 256
modules = []
modules.append(nn.Sequential(
 nn.Conv2d(in_channels, out_channels, 1, bias=False),
nn.BatchNorm2d(out_channels),
nn.ReLU(inplace=True)))
 rate1, rate2, rate3 = tuple(atrous_rates)
 modules.append(ASPPConv(in_channels, out_channels, rate1))
modules.append(ASPPConv(in_channels, out_channels, rate2))
modules.append(ASPPConv(in_channels, out_channels, rate3))
modules.append(ASPPPooling(in_channels, out_channels))
self.convs = nn.ModuleList(modules)
 self.project = nn.Sequential(
 nn.Conv2d(5 * out_channels, out_channels, 1, bias=False),
nn.BatchNorm2d(out_channels),
nn.ReLU(inplace=True),
nn.Dropout(0.1),)
 def forward(self, x):
res = []
for conv in self.convs:
res.append(conv(x))
res = torch.cat(res, dim=1)
 return self.project(res )
model.eval()
prev_test_df = None
MAX_SEQ = 100
for(test_df, sample_prediction_df)in tqdm(iter_test):
if(prev_test_df is not None)&(psutil.virtual_memory().percent<90):
print(psutil.virtual_memory().percent)
 prev_test_df['answered_correctly'] = eval(test_df['prior_group_answers_correct'].iloc[0])
 prev_test_df = prev_test_df[prev_test_df.content_type_id == False]
prev_group = prev_test_df[['user_id', 'content_id', 'answered_correctly']].groupby('user_id' ).apply(lambda r:(
 r['content_id'].values,
r['answered_correctly'].values))
for prev_user_id in prev_group.index:
prev_group_content = prev_group[prev_user_id][0]
prev_group_ac = prev_group[prev_user_id][1]
if prev_user_id in group.index:
group[prev_user_id] =(np.append(group[prev_user_id][0],prev_group_content),
np.append(group[prev_user_id][1],prev_group_ac))
else:
group[prev_user_id] =(prev_group_content,prev_group_ac)
 if len(group[prev_user_id][0])>MAX_SEQ:
new_group_content = group[prev_user_id][0][-MAX_SEQ:]
new_group_ac = group[prev_user_id][1][-MAX_SEQ:]
group[prev_user_id] =(new_group_content,new_group_ac)
 prev_test_df = test_df.copy()
test_df = test_df[test_df.content_type_id == False]
test_dataset = TestDataset(group, test_df, skills)
 test_dataloader = DataLoader(test_dataset, batch_size=51200, shuffle=False)
 outs = []
for item in tqdm(test_dataloader):
x = item[0].to(device ).long()
target_id = item[1].to(device ).long()
with torch.no_grad() :
output, att_weight = model(x, target_id)
 output = torch.sigmoid(output)
 output = output[:, -1]
outs.extend(output.view(-1 ).data.cpu().numpy())
 test_df['answered_correctly'] = outs
env.predict(test_df.loc[test_df['content_type_id'] == 0, ['row_id', 'answered_correctly']] )<install_modules>
def mask_to_rle(mask):

 pixels= mask.flatten()
pixels = np.concatenate([[0], pixels, [0]])
 runs = np.where(pixels[1:] != pixels[:-1])[0] + 1
runs[1::2] -= runs[::2]
return ' '.join(str(x)for x in runs )
!pip install -q efficientnet >> /dev/null
Conv2D,
BatchNormalization,
Dense,
GlobalAveragePooling2D,
Input,
Activation,
Lambda,
multiply)
 sns.set_style("darkgrid" )<define_variables>
from torch.utils.data import DataLoader
from tqdm import tqdm
import cv2
DEVICE = "TPU"
SEED = 42
FOLDS = 5
IMG_SIZES = [384,384,384,384,384]
INC2019 = [0,0,0,0,0]
INC2018 = [1,1,1,1,1]
BATCH_SIZES = [32]*FOLDS
EPOCHS = [12]*FOLDS
EFF_NETS = [6,6,6,6,6]
WGTS = [1/FOLDS]*FOLDS
TTA = 11<choose_model_class>
!cp -r.. /input/pytorch-segmentation-models/segmentation_models_pytorch-0.1.3-py3-none-any.whl /kaggle/working
!cp -r.. /input/pytorch-segmentation-models/timm-0.3.2-py3-none-any.whl /kaggle/working
!cp -r.. /input/smpwhl/efficientnet_pytorch-0.6.3-py3-none-any.whl /kaggle/working
!cp -r.. /input/smpwhl/pretrainedmodels-0.7.4-py3-none-any.whl /kaggle/working
!pip install /kaggle/working/timm-0.3.2-py3-none-any.whl
!pip install /kaggle/working/efficientnet_pytorch-0.6.3-py3-none-any.whl
!pip install /kaggle/working/pretrainedmodels-0.7.4-py3-none-any.whl
!pip install /kaggle/working/segmentation_models_pytorch-0.1.3-py3-none-any.whl
if DEVICE == "TPU":
print("connecting to TPU...")
 try:
tpu = tf.distribute.cluster_resolver.TPUClusterResolver()
print('Running on TPU ', tpu.master())
 except ValueError:
print("Could not connect to TPU")
 tpu = None
if tpu:
try:
print("initializing TPU...")
 tf.config.experimental_connect_to_cluster(tpu)
 tf.tpu.experimental.initialize_tpu_system(tpu)
 strategy = tf.distribute.experimental.TPUStrategy(tpu)
 print("TPU initialized")
 except _:
print("failed to initialize TPU")
 else:
DEVICE = "GPU"
if DEVICE != "TPU":
print("Using default strategy for CPU and single GPU")
 strategy = tf.distribute.get_strategy()
if DEVICE == "GPU":
print("Num GPUs Available: ", len(tf.config.experimental.list_physical_devices('GPU')))
 AUTO = tf.data.experimental.AUTOTUNE
REPLICAS = strategy.num_replicas_in_sync
print(f'REPLICAS: {REPLICAS}' )<feature_engineering>
model = smp.Unet(
 encoder_name="timm-efficientnet-b5",
encoder_weights = None,
in_channels=3,
classes=5, 
)
 model_path = '.. /input/u-enetb5-full-all-ep110/enb5_net_dice96.04_ep110.pth'
model = load_model(model, model_path)
 model = model.cuda()
GCS_PATH = [None]*FOLDS; GCS_PATH2 = [None]*FOLDS
for i,k in enumerate(IMG_SIZES):
GCS_PATH[i] = KaggleDatasets().get_gcs_path('melanoma-%ix%i'%(k,k))
GCS_PATH2[i] = KaggleDatasets().get_gcs_path('isic2019-%ix%i'%(k,k))
files_train = np.sort(np.array(tf.io.gfile.glob(GCS_PATH[0] + '/train*.tfrec')))
 files_test = np.sort(np.array(tf.io.gfile.glob(GCS_PATH[0] + '/test*.tfrec')) )<define_variables>
defect_root = '/kaggle/input/severstal-steel-defect-detection/test_images'
dataset = SteelDataset(defect_root, resize_long=(800, 800), img_normal=True, to_torch=True)
 data_loader = DataLoader(dataset, batch_size=4, shuffle=False, num_workers=4, pin_memory=True )
<normalization><EOS> 
device='cuda:0' if torch.cuda.is_available() else 'cpu'
image_id_class_id = []
encoded_pixel = []
threshold_label = [ 0.75, 0.85, 0.50, 0.50,]
threshold_mask_pixel = [ 0.40, 0.40, 0.40, 0.40,]
threshold_mask_size = [ 50, 250, 1400, 1800,]
with torch.no_grad() :
for i_iter, batch in enumerate(data_loader):
img_name = [name.split('/')[-1] for name in batch['name']]
img = batch['img'].to(device)
 if 1:
logit = model(img)
 if isinstance(logit, list):
logit = logit[1]
logit = F.interpolate(logit,(256, 1600), mode='bilinear', align_corners=True)
 probability = torch.softmax(logit,1)
 probability_mask = probability[:,1:]
probability_label = probability_mask_to_probability_label(probability)[:,1:]
num_augment = 1
if 1 :
logit = model(torch.flip(img,dims=[3]))
if isinstance(logit, list):
logit = logit[1]
logit = F.interpolate(logit,(256, 1600), mode='bilinear', align_corners=True)
 probability = torch.softmax(torch.flip(logit,dims=[3]),1)
 probability_mask += probability[:,1:]
probability_label += probability_mask_to_probability_label(probability)[:,1:]
num_augment +=1
if 1 :
logit = model(torch.flip(img,dims=[2]))
if isinstance(logit, list):
logit = logit[1]
logit = F.interpolate(logit,(256, 1600), mode='bilinear', align_corners=True)
 probability = torch.softmax(torch.flip(logit,dims=[2]),1)
 probability_mask += probability[:,1:]
probability_label += probability_mask_to_probability_label(probability)[:,1:]
num_augment +=1
probability_mask = probability_mask/num_augment
probability_label = probability_label/num_augment
probability_mask = probability_mask.data.cpu().numpy()
probability_label = probability_label.data.cpu().numpy()
for b in range(img.shape[0]):
for c in range(4):
rle=''
predict_label = probability_label[b,c]>threshold_label[c]
if predict_label:
try:
predict_mask = probability_mask[b,c] > threshold_mask_pixel[c]
predict_mask = post_process(predict_mask, threshold_mask_size[c])
 rle = run_length_encode(predict_mask)
 except:
print('An exception occurred : %s'%(image_id[b]+'_%d'%(c+1)))
 image_id_class_id.append(img_name[b]+'_%d'%(c+1))
encoded_pixel.append(rle)
 if(i_iter+1)% 100 == 0:
print('{}/{}'.format(i_iter+1,len(dataset)) , 'time:', time.time() -s_time)
 df = pd.DataFrame(zip(image_id_class_id, encoded_pixel), columns=['ImageId_ClassId', 'EncodedPixels'])
 df.to_csv('submission.csv', index=False)
 print('total_time:', time.time() -s_time )